diff --git a/examples/moe_recipes/README.md b/examples/moe_recipes/README.md new file mode 100644 index 00000000000..a39a3738c6f --- /dev/null +++ b/examples/moe_recipes/README.md @@ -0,0 +1,14 @@ +# MoE Recipes + +Each recipe contains a self-contained container definition, runtime environment +variables, and Megatron-LM arguments. + +The DeepSeek V4 recipe is launched with `pretrain_hybrid.py`. Set `LOAD_PATH` +and `OUTPUT_PATH` before rendering its arguments. + +| Model | Recipe | GPUs | TP/PP/EP/CP/ETP | MBS/GBS/SL | Features | +|---|---|---:|---|---|---| +| DeepSeek-V4-Flash | [GB200 MXFP8 THD 64K](deepseek_v4_flash/gb200/mxfp8_THD_SL64K_128GPU_TP1PP2EP64CP16.yaml) | 128 | 1/2/64/16/1 | 1/128/65536 | Hybrid attention; THD packing; mHC; MTP; HybridEP; scoped TE graphs; fine-grained activation offload | + +TP = tensor parallel, PP = pipeline parallel, EP = expert parallel, CP = +context parallel, and ETP = expert tensor parallel. diff --git a/examples/moe_recipes/deepseek_v4_flash/gb200/mxfp8_THD_SL64K_128GPU_TP1PP2EP64CP16.yaml b/examples/moe_recipes/deepseek_v4_flash/gb200/mxfp8_THD_SL64K_128GPU_TP1PP2EP64CP16.yaml new file mode 100644 index 00000000000..9b90029b6e1 --- /dev/null +++ b/examples/moe_recipes/deepseek_v4_flash/gb200/mxfp8_THD_SL64K_128GPU_TP1PP2EP64CP16.yaml @@ -0,0 +1,235 @@ +# Launch this recipe's ARGS with pretrain_hybrid.py. +DEPENDENCIES: + pytorch_base_image: nvcr.io/nvidia/pytorch:26.04-py3 + dockerfile: | + # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + # IMAGE_NAME: dsv4-hybrid-gb200-torch2604 + + FROM nvcr.io/nvidia/pytorch:26.04-py3 + + ENV SHELL=/bin/bash + ENV DEBIAN_FRONTEND=noninteractive + + RUN apt-get update && \ + apt-get install -y --no-install-recommends git curl wget gettext sudo && \ + wget https://github.com/mikefarah/yq/releases/download/v4.27.5/yq_linux_arm64 \ + -O /usr/bin/yq && \ + chmod +x /usr/bin/yq && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* + + RUN unset PIP_CONSTRAINT && \ + pip install --no-cache-dir \ + datasets einops omegaconf sentencepiece tensorboard tiktoken \ + tokenizers transformers==4.57.1 wandb \ + nvidia-cutlass-dsl==4.5.2 nvidia-cudnn-frontend==1.26.0 + + # Pins match the GB200 image used to validate this port. + ARG TE_COMMIT=7cb8b313d55021eb12c5efd4d02f4c8ff7679453 + RUN unset PIP_CONSTRAINT && \ + NVTE_CUDA_ARCHS="100a;103a" NVTE_BUILD_THREADS_PER_JOB=8 \ + NVTE_FRAMEWORK=pytorch pip install --no-build-isolation --no-cache-dir \ + "git+https://github.com/NVIDIA/TransformerEngine.git@$TE_COMMIT" + + ARG HYBRIDEP_COMMIT=1b8f467965bb818bf2f6511e06993f5607e1721f + RUN git clone --branch hybrid-ep https://github.com/deepseek-ai/DeepEP.git \ + /workspace/DeepEP && \ + cd /workspace/DeepEP && \ + git checkout $HYBRIDEP_COMMIT && \ + TORCH_CUDA_ARCH_LIST="10.0" pip install --no-build-isolation . + + ARG FHT_COMMIT=e7706faf8d1c3b9f241e36860640ad1dac644ede + RUN git clone https://github.com/Dao-AILab/fast-hadamard-transform.git \ + /workspace/fast-hadamard-transform && \ + cd /workspace/fast-hadamard-transform && \ + git checkout $FHT_COMMIT && \ + pip install --no-build-isolation . + + ARG FLASHMLA_COMMIT=b7643bd54521f563b839b98289b5cd048c062ba2 + RUN git clone --branch nv_dev https://github.com/deepseek-ai/FlashMLA.git \ + /workspace/FlashMLA && \ + cd /workspace/FlashMLA && \ + git checkout $FLASHMLA_COMMIT && \ + FLASH_MLA_DISABLE_SM90=1 NVCC_THREADS=16 \ + CFLAGS="-I$CUDA_HOME/include/cccl" \ + CXXFLAGS="-I$CUDA_HOME/include/cccl" \ + pip install --no-build-isolation . + + RUN rm -rf /root/.cache /tmp/* + WORKDIR /workspace + +ENV_VARS: + TORCH_NCCL_AVOID_RECORD_STREAMS: '0' + NVTE_ALLOW_NONDETERMINISTIC_ALGO: '1' + PYTORCH_CUDA_ALLOC_CONF: expandable_segments:True,graph_capture_record_stream_reuse:True + NCCL_NVLS_ENABLE: '0' + NVTE_FUSED_ATTN: '1' + NVTE_NORM_FWD_USE_CUDNN: '1' + NVTE_NORM_BWD_USE_CUDNN: '1' + PYTHONWARNINGS: ignore + NCCL_DEBUG: VERSION + NCCL_GRAPH_REGISTER: '0' + NVTE_CUTEDSL_FUSED_GROUPED_MLP: '1' + NVTE_CPU_OFFLOAD_V1: '1' + NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN: '64' + USE_MNNVL: '1' + NUM_OF_TOKENS_PER_CHUNK_COMBINE_API: '128' + NUM_OF_STAGES_DISPATCH_API: '10' + NUM_OF_IN_FLIGHT_S2G_DISPATCH_API: '8' + +ARGS: + tokenizer_type: HuggingFaceTokenizer + tokenizer_model: unsloth/DeepSeek-V3 + hidden_size: 4096 + num_attention_heads: 64 + kv_channels: 512 + max_position_embeddings: 65536 + normalization: RMSNorm + norm_epsilon: 1e-6 + swiglu: true + disable_bias_linear: true + untie_embeddings_and_output_weights: true + position_embedding_type: rope + rotary_base: 10000 + make_vocab_size_divisible_by: 3232 + multi_latent_attention: true + q_lora_rank: 1024 + qk_pos_emb_head_dim: 64 + v_head_dim: 512 + rotary_scaling_factor: 4 + mscale: 1.0 + mscale_all_dim: 1.0 + qk_layernorm: true + o_groups: 8 + o_lora_rank: 1024 + original_max_position_embeddings: 65536 + experimental_attention_variant: dsv4_hybrid + csa_window_size: 128 + csa_compress_ratios: '([0,0,4]+[128,4]*20+[0])' + csa_compress_rotary_base: 40000 + dsa_indexer_n_heads: 64 + dsa_indexer_head_dim: 128 + dsa_indexer_topk: 512 + dsa_indexer_loss_coeff: 1e-2 + dsa_indexer_use_sparse_loss: true + num_experts: 256 + moe_n_hash_layers: 3 + moe_ffn_hidden_size: 2048 + moe_shared_expert_intermediate_size: 2048 + moe_router_load_balancing_type: seq_aux_loss + moe_router_topk: 6 + moe_aux_loss_coeff: 1e-4 + moe_router_topk_scaling_factor: 1.5 + moe_router_score_function: sqrtsoftplus + moe_router_enable_expert_bias: true + moe_router_bias_update_rate: 1e-3 + activation_func_clamp_value: 10.0 + enable_hyper_connections: true + num_residual_streams: 4 + mhc_sinkhorn_iterations: 20 + use_fused_mhc: true + mtp_num_layers: 1 + mtp_loss_scaling_factor: 0.1 + attention_dropout: 0.0 + hidden_dropout: 0.0 + hybrid_layer_pattern: "WEWECEHECE|HECEHECEHECE|HECEHECEHECE|HECEHECEHECE|HECEHECEHECE|HECEHECEHECE|HECEHECEHECE|HECE/WE" + spec: + - megatron.core.models.hybrid.hybrid_layer_specs + - hybrid_dsv4_stack_spec + mock_data: true + seq_length: 65536 + moe_router_force_load_balancing: true + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 2 + expert_model_parallel_size: 64 + context_parallel_size: 16 + expert_tensor_parallel_size: 1 + use_distributed_optimizer: true + overlap_grad_reduce: true + overlap_param_gather: true + moe_token_dispatcher_type: flex + moe_flex_dispatcher_backend: hybridep + moe_hybridep_num_sms: 32 + moe_grouped_gemm: true + moe_permute_fusion: true + moe_router_fusion: true + moe_router_dtype: fp32 + recompute_granularity: selective + recompute_modules: + - mla_up_proj + fine_grained_activation_offloading: true + offload_modules: + - expert_fc1 + delay_offload_until_cuda_graph: true + cuda_graph_impl: transformer_engine + cuda_graph_modules: + - attn + - moe_router + - moe_preprocess + te_rng_tracker: true + cuda_graph_warmup_steps: 1 + use_flash_attn: true + transformer_impl: transformer_engine + micro_batch_size: 1 + global_batch_size: 128 + train_samples: 585937500 + exit_duration_in_mins: 220 + no_save_optim: true + no_check_for_nan_in_loss_and_grad: true + cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: native + no_create_attention_mask_in_dataloader: true + manual_gc: true + manual_gc_interval: 10 + use_varlen_dataset: true + varlen_mock_dataset_config_json: '{"mode":"distribution","type":"lognormal","format":"thd","min_seq_len":65536,"max_seq_len":65536,"mean_seq_len":65536,"lognormal_sigma":1.1}' + sequence_packing_scheduler: dp_balanced + calculate_per_token_loss: true + pad_packed_seq_alignment: max + max_seqlen_per_dp_cp_rank: 4096 + thd_max_packed_sequences: 8 + cp_partition_mode: contiguous + lr: 3.9e-06 + min_lr: 3.9e-07 + lr_warmup_init: 3.9e-07 + lr_decay_style: cosine + lr_decay_samples: 584765624 + lr_warmup_samples: 1536000 + weight_decay: 0.1 + clip_grad: 1.0 + adam_beta1: 0.9 + adam_beta2: 0.95 + bf16: true + fp8_recipe: mxfp8 + fp8_format: e4m3 + fp8_param_gather: true + reuse_grad_buf_for_mxfp8_param_ag: true + use_precision_aware_optimizer: true + main_grads_dtype: fp32 + main_params_dtype: fp32 + exp_avg_dtype: bf16 + exp_avg_sq_dtype: bf16 + moe_router_padding_for_quantization: true + init_method_std: 0.02 + eval_iters: 32 + eval_interval: 200 + no_load_optim: true + no_load_rng: true + auto_detect_ckpt_format: true + load: $LOAD_PATH + save: $OUTPUT_PATH/checkpoints + save_interval: 500 + dist_ckpt_strictness: log_all + log_throughput: true + log_interval: 1 + logging_level: 20 + log_timers_to_tensorboard: true + log_memory_to_tensorboard: true + log_validation_ppl_to_tensorboard: true + tensorboard_dir: $OUTPUT_PATH/tensorboard + wandb_exp_name: DeepSeek-V4-Flash-HybridModel-GB200-MXFP8-THD-TP1PP2EP64CP16-GBS128SEQLEN65536 + wandb_project: megatron-deepseek-v4-flash-benchmark + wandb_save_dir: $OUTPUT_PATH/wandb + enable_experimental: true + log_memory_interval: 5 + log_device_memory_used: true diff --git a/hybrid_builders.py b/hybrid_builders.py index 7e1c58682ac..4595f904c17 100644 --- a/hybrid_builders.py +++ b/hybrid_builders.py @@ -1,12 +1,12 @@ # Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. -from model_provider import count_parameters_in_layer +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_inference_stack_spec from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.transformer import TransformerConfig -from megatron.core.transformer.spec_utils import import_module +from megatron.core.transformer.spec_utils import ModuleSpec, import_module from megatron.training import print_rank_0 from megatron.training.arguments import core_transformer_config_from_args -from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_inference_stack_spec +from model_provider import count_parameters_in_layer def hybrid_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_collection=None): @@ -21,6 +21,8 @@ def hybrid_builder(args, pre_process, post_process, vp_stage=None, config=None, ), "inference_fuse_tp_communication is not supported for HybridModel" elif args.spec is not None: hybrid_stack_spec = import_module(args.spec) + if not isinstance(hybrid_stack_spec, ModuleSpec) and callable(hybrid_stack_spec): + hybrid_stack_spec = hybrid_stack_spec(config) else: raise ValueError("You must provide a valid hybrid layer spec via --spec") diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index 0f016473b6a..23eae2799b8 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -1,12 +1,69 @@ # Copyright (c) 2025 NVIDIA CORPORATION. All rights reserved. -from typing import Any, List, Optional +from typing import Any, Dict, List, Optional, Type import torch from megatron.core import parallel_state +from megatron.core.datasets.data_schedule_utils import ( + broadcast_scalars, + broadcast_tensor, + build_packed_microbatches, + create_data_iterator, + get_batch_and_global_seqlens, + get_cp_slice_for_thd, + reroute_samples_to_dcp_ranks, +) +from megatron.core.packed_seq_params import ( + PackedSeqParams, + get_thd_padding_kwargs, + pad_sequence_for_thd, +) from megatron.core.pipeline_parallel.hybrid_cp_schedule import BalancedCPScheduler from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.multi_token_prediction import mtp_on_this_rank as mtp_is_on_rank + + +def _build_thd_padding_mask( + cu_seqlens: torch.Tensor, cu_seqlens_padded: torch.Tensor +) -> torch.Tensor: + """Build a 1D THD padding mask from scheduler sequence metadata.""" + assert cu_seqlens.dim() == 1 + assert cu_seqlens_padded.dim() == 1 + assert cu_seqlens.numel() == cu_seqlens_padded.numel() + + total_tokens = int(cu_seqlens_padded[-1].item()) + if total_tokens == 0: + return torch.empty((0,), dtype=torch.bool, device=cu_seqlens.device) + + num_sequences = cu_seqlens.numel() - 1 + if num_sequences <= 0: + return torch.ones((total_tokens,), dtype=torch.bool, device=cu_seqlens.device) + + positions = torch.arange( + total_tokens, dtype=cu_seqlens_padded.dtype, device=cu_seqlens_padded.device + ) + seq_indices = torch.searchsorted(cu_seqlens_padded[1:].contiguous(), positions, right=True) + + valid_lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).clamp(min=0) + valid_ends = cu_seqlens_padded[:-1] + valid_lengths + return positions >= valid_ends[seq_indices] + + +def _sanitize_thd_padding_values(batch: Dict[str, Any], padding_mask: torch.Tensor) -> None: + """Replace padded token-like slots with safe neutral values in-place.""" + assert padding_mask.dim() == 1 + pad_values = {'tokens': 0, 'labels': 0, 'loss_mask': 0.0, 'position_ids': 0} + for key, pad_value in pad_values.items(): + tensor = batch.get(key) + if tensor is None: + continue + assert tensor.dim() == 1, f"{key} must be 1D before CP slicing, got {tensor.dim()}D" + assert tensor.numel() == padding_mask.numel(), ( + f"{key} length ({tensor.numel()}) must match padding_mask length " + f"({padding_mask.numel()}) before CP slicing." + ) + batch[key] = tensor.masked_fill(padding_mask, pad_value) class HybridCPDataLoaderWrapper: @@ -57,7 +114,6 @@ def get_global_seqlens(self, subsample_seqlens: torch.Tensor) -> List[int]: Gathers the sequence lengths of all subsamples from all DP ranks. Each DP rank loads the same number of microbatches but each microbatch may have a different number of subsamples. - We find the number of subsamples each rank holds and then gather the sequence lengths of all subsamples from all ranks. """ @@ -299,3 +355,651 @@ def __next__(self) -> Any: batch, global_ids_this_rank, global_id_seqlens, sample_id_groups, offsets ) return samples_this_rank_with_id, sample_id_groups + + +class BasePackingScheduler: + """Base class for sequence packing schedulers.""" + + def __init__( + self, + max_seqlen_per_dp_cp_rank: int, + cp_size: int, + dp_size: int, + microbatch_group_size_per_vp_stage: Optional[int], + max_num_seqs: Optional[int] = None, + ): + """ + Args: + max_seqlen_per_dp_cp_rank: The maximum sequence length per DPxCP rank. + cp_size: The context parallel size. + dp_size: The data parallel size. + microbatch_group_size_per_vp_stage: The microbatch group size per virtual + pipeline stage, only used when enabling VPP, otherwise None. + """ + self.max_seqlen_per_dp_cp_rank = max_seqlen_per_dp_cp_rank + self.cp_size = cp_size + self.dp_size = dp_size + self.microbatch_group_size_per_vp_stage = microbatch_group_size_per_vp_stage + self.max_num_seqs = max_num_seqs + + def get_required_sample_keys(self): + """Return the required key of each batch.""" + raise NotImplementedError + + def get_groups_and_subsamples(self, sample_id_seqlens): + """schedule the samples into groups""" + raise NotImplementedError + + def run( + self, + data_iterator, + num_microbatches, + dp_group, + tp_group, + pp_group, + dp_cp_group, + dev, + config, + ): + """ + Run the scheduler and return the new data_iterator. + + Args: + data_iterator: The data iterator. + num_microbatches: The number of microbatches to fetch. + dp_group: Data parallel process group. + tp_group: Tensor parallel process group. + pp_group: Pipeline parallel process group. + dp_cp_group: Data parallel + context parallel process group. + dev: CUDA device. + config: Model parallel config. + + Returns: + new_data_iterator: The new data iterator (or list for VPP). + num_micro_batches: Number of micro batches after scheduling. + seqlen_sum_this_global_batch: Total tokens for FLOPs calculation. + seqlen_squared_sum_this_global_batch: Sum of squared seqlens for FLOPs. + """ + raise NotImplementedError + + +class DpBalancedScheduler(BasePackingScheduler): + """Packs sequences in their original order until reaching the max limit of sequence length.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.max_seq_len_all_ranks = self.max_seqlen_per_dp_cp_rank * self.cp_size + + def get_required_sample_keys(self): + """Return the required key of each batch.""" + return [ + "tokens", + "labels", + "loss_mask", + "position_ids", + "original_seq_len", # Length of the original sequence length, should be a gpu tensor. + "padded_seq_len", # Length of the padded sequence length, should be a gpu tensor. + ] + + def get_groups_and_subsamples(self, sample_id_seqlens): + """ + Packs sequences in their original order until reaching the max limit of sequence length. + """ + sample_id_groups = [] + packed_id_groups = [] + sum_seqlen = 0 + single_microbatch = [] + + for i in range(len(sample_id_seqlens)): + if sum_seqlen + sample_id_seqlens[i][1] <= self.max_seq_len_all_ranks and ( + self.max_num_seqs is None or len(single_microbatch) < self.max_num_seqs + ): + single_microbatch.append(i) + sum_seqlen += sample_id_seqlens[i][1] + else: + packed_id_groups.append(single_microbatch) + single_microbatch = [i] + sum_seqlen = sample_id_seqlens[i][1] + if len(single_microbatch) > 0: + packed_id_groups.append(single_microbatch) + + # we want the number of packed sequences to be multiple of dp_size + # so we move few samples from previous microbatch + # to the end of the microbatches if needed + num_packed_sequence = len(packed_id_groups) + + # when enabling vpp, we want the number of packed sequences to be + # multiple of dp_size * microbatch_group_size_per_vp_stage + multiple = self.dp_size * ( + self.microbatch_group_size_per_vp_stage + if self.microbatch_group_size_per_vp_stage is not None + else 1 + ) + if num_packed_sequence % multiple != 0: + remainder = num_packed_sequence % multiple + num_to_move = multiple - remainder + i = num_packed_sequence - 1 + while num_to_move > 0: + assert i >= 0, "Not enough samples to move" + if len(packed_id_groups[i]) > 1: + seq_id = packed_id_groups[i].pop() + packed_id_groups.append([seq_id]) + num_to_move -= 1 + else: + i -= 1 + + num_micro_batches = int(len(packed_id_groups) / self.dp_size) + for i in range(num_micro_batches): + sample_id_groups.append([]) + for j in range(self.cp_size * self.dp_size): + seq_id = int(i * self.dp_size + j / self.cp_size) + sample_id_groups[i].append(packed_id_groups[seq_id]) + return sample_id_groups + + def run( + self, + data_iterator, + num_microbatches: int, + dp_group, + tp_group, + pp_group, + dp_cp_group, + dev: torch.device, + config, + ): + """ + Run the complete scheduling pipeline. + + Packed-sequence datasets are built on TP rank 0 of every PP stage. Each + stage therefore runs the same schedule locally, retaining only the data + fields required by that stage before the all-to-all transfer. + + Args: + data_iterator: The data iterator. + num_microbatches: The number of microbatches to fetch. + dp_group: Data parallel process group. + tp_group: Tensor parallel process group. + pp_group: Pipeline parallel process group. + dp_cp_group: Data parallel + context parallel process group. + dev: CUDA device. + config: Model parallel config. + + Returns: + new_data_iterator: The new data iterator (or list for VPP). + num_micro_batches: Number of micro batches after scheduling. + seqlen_sum_this_global_batch: Total tokens for FLOPs calculation. + seqlen_squared_sum_this_global_batch: Sum of squared seqlens for FLOPs. + """ + + total_dcp_gpus = dp_cp_group.size() + is_first_pp = pp_group.rank() == 0 + is_last_pp = pp_group.rank() == pp_group.size() - 1 + mtp_on_this_pp = mtp_is_on_rank( + layout=config.pipeline_model_parallel_layout, + mtp_num_layers=config.mtp_num_layers, + ignore_virtual=True, + ) + + vpp_size = config.virtual_pipeline_model_parallel_size or 1 + vpp_needs_data = None + if vpp_size > 1: + assert len(data_iterator) == vpp_size + data_iterator = next( + (iterator for iterator in data_iterator if iterator is not None), None + ) + + vpp_needs_data = [False] * vpp_size + if is_first_pp: + vpp_needs_data[0] = True + if is_last_pp: + vpp_needs_data[-1] = True + if mtp_on_this_pp: + for vp_stage in range(vpp_size): + if mtp_is_on_rank( + layout=config.pipeline_model_parallel_layout, + mtp_num_layers=config.mtp_num_layers, + ignore_virtual=False, + vp_stage=vp_stage, + ): + vpp_needs_data[vp_stage] = True + + if data_iterator is not None: + assert tp_group.rank() == 0, "Only TP rank 0 should have data_iterator" + + # Step 1: Fetch batches and gather global sequence lengths + ( + batch, + global_id_seqlens, + global_ids_this_rank, + offsets, + _padded_seqlens_gathered, + original_seqlens_gathered, + ) = get_batch_and_global_seqlens(data_iterator, num_microbatches, dp_group) + + # Step 2: Check required sample keys + for key in self.get_required_sample_keys(): + assert ( + key in batch[0] + ), f"Batch missing required key {key}, provided keys: {batch[0].keys()}" + + # Avoid transferring fields that this pipeline stage never consumes. + keys_to_keep = {'original_seq_len', 'padded_seq_len'} + if is_first_pp or mtp_on_this_pp: + keys_to_keep.update(['tokens', 'position_ids']) + if is_last_pp or mtp_on_this_pp: + keys_to_keep.update(['labels', 'loss_mask']) + for sample in batch: + for key in list(sample): + if key not in keys_to_keep: + del sample[key] + + # Step 3: Schedule samples into groups + sample_id_groups = self.get_groups_and_subsamples(global_id_seqlens) + + # Validate scheduling result + set_gbs = set() + for group in sample_id_groups: + for sub in group: + set_gbs.update(sub) + assert len(set_gbs) == len(global_id_seqlens), ( + f"set_gbs length: {len(set_gbs)} != " + f"global_id_seqlens length: {len(global_id_seqlens)}" + ) + + # Step 4: Reroute samples to DCP ranks + samples_this_rank_with_id = reroute_samples_to_dcp_ranks( + batch, + global_ids_this_rank, + global_id_seqlens, + sample_id_groups, + offsets, + dp_group, + dp_cp_group, + total_dcp_gpus, + ) + + dcp_rank = dp_cp_group.rank() + num_micro_batches = len(sample_id_groups) + + grouped_samples = [ + [ + samples_this_rank_with_id[sub_sample_id] + for sub_sample_id in sample_id_groups[i][dcp_rank] + ] + for i in range(num_micro_batches) + ] + + # Step 5: Build packed microbatches + new_samples = build_packed_microbatches(grouped_samples, dev) + + # Step 6: Calculate FLOPs info + seqlen_sum_this_global_batch = float(sum(original_seqlens_gathered)) + seqlen_squared_sum_this_global_batch = float( + sum(seqlen**2 for seqlen in original_seqlens_gathered) + ) + else: + ( + new_samples, + num_micro_batches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ) = (None, None, None, None) + + # Broadcast scalar schedule results to the remaining TP ranks. + (num_micro_batches, seqlen_sum_this_global_batch, seqlen_squared_sum_this_global_batch) = ( + broadcast_scalars( + [ + num_micro_batches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ], + tp_group, + dev, + ) + ) + num_micro_batches = int(num_micro_batches) + + new_data_iterator = create_data_iterator(new_samples, tp_group, config, vpp_needs_data) + + return ( + new_data_iterator, + num_micro_batches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ) + + +scheduler_map: Dict[str, Type[BasePackingScheduler]] = {'dp_balanced': DpBalancedScheduler} + + +def _get_scheduler_max_real_num_seqs(config) -> Optional[int]: + """Return the real-sequence cap after reserving an optional dummy-tail slot.""" + max_num_seqs = getattr(config, 'thd_max_packed_sequences', None) + if max_num_seqs is None: + return None + + max_num_seqs = int(max_num_seqs) + if max_num_seqs < 1: + raise ValueError(f"thd_max_packed_sequences must be >= 1, got {max_num_seqs}.") + if getattr(config, 'pad_packed_seq_alignment', None) is not None and getattr( + config, 'pad_packed_seq_by_appending_dummy_seq', True + ): + if max_num_seqs < 2: + raise ValueError( + "thd_max_packed_sequences must be >= 2 when THD padding appends " + "a dummy sequence." + ) + return max_num_seqs - 1 + return max_num_seqs + + +def wrap_data_iterator( + data_iterator, + config, + num_microbatches, + pg_collection: ProcessGroupCollection, +): + """ + A wrapper function that wraps around an existing data_iterator + and return the num_micro_batches for sequence packing. + + Args: + data_iterator: The original data_iterator to wrap around + config: The config object containing the max_seqlen_per_dp_cp_rank + pg_collection: The process group collection. + """ + dp_cp_group = pg_collection.dp_cp + dp_group = pg_collection.dp + tp_group = pg_collection.tp + pp_group = pg_collection.pp + assert ( + dp_cp_group is not None + and dp_group is not None + and tp_group is not None + and pp_group is not None + ), "dp_cp_group, dp_group, tp_group must not be None when using sequence packing" + + dev = torch.cuda.current_device() + dp_size = dp_group.size() + cp_size = dp_cp_group.size() // dp_size + + scheduler_type = config.sequence_packing_scheduler + scheduler_max_num_seqs = ( + _get_scheduler_max_real_num_seqs(config) + if scheduler_type == 'dp_balanced' + else getattr(config, 'thd_max_packed_sequences', None) + ) + scheduler = scheduler_map[scheduler_type]( + config.max_seqlen_per_dp_cp_rank, + cp_size, + dp_size, + # When VPP is enabled, align num_micro_batches to this multiple. + ( + None + if config.virtual_pipeline_model_parallel_size is None + else config.microbatch_group_size_per_vp_stage + ), + max_num_seqs=scheduler_max_num_seqs, + ) + + ( + new_data_iterator, + num_micro_batches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ) = scheduler.run( + data_iterator, num_microbatches, dp_group, tp_group, pp_group, dp_cp_group, dev, config + ) + + return ( + new_data_iterator, + num_micro_batches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ) + + +def get_batch_on_this_rank_for_sequence_packing( + data_iterator, + pg_collection: ProcessGroupCollection, + vpp_size: Optional[int] = None, + mtp_on_this_rank: bool = False, + vp_stage: Optional[int] = None, + config=None, +): + """ + Get a batch of data for sequence packing. + Args: + 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. + pg_collection: The process group collection. + 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) + """ + + tp_group = pg_collection.tp + pp_group = pg_collection.pp + cp_group = pg_collection.cp + + tp_src_rank = torch.distributed.get_process_group_ranks(tp_group)[0] + + is_tp_rank_0 = tp_group.rank() == 0 + is_first_stage = pp_group.rank() == 0 and (vp_stage is None or vp_stage == 0) + is_last_stage = pp_group.rank() == pp_group.size() - 1 and ( + vp_stage is None or vp_stage == vpp_size - 1 + ) + + dev = torch.cuda.current_device() + + # data_iterator should return a batch including the following keys. + batch_keys = ['cu_seqlens', 'cu_seqlens_padded', 'max_seqlen'] + if is_first_stage or mtp_on_this_rank: + batch_keys.append('tokens') + batch_keys.append('position_ids') + if is_last_stage or mtp_on_this_rank: + batch_keys.append('labels') + batch_keys.append('loss_mask') + + # Get a batch from data_iterator or create an emtpy batch. + if is_tp_rank_0: + assert data_iterator is not None + batch = next(data_iterator) + for key in batch_keys: + assert key in batch, f"{key} is missing in current batch." + else: + assert data_iterator is None, "Non TP 0 rank should not have data_iterator" + batch = {} + + 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') == 'transformer_engine', + ) + 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 + # the origin assigned to each CP rank. + 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: + batch['padding_mask'] = _build_thd_padding_mask( + batch['cu_seqlens'], batch['cu_seqlens_padded'] + ) + _sanitize_thd_padding_values(batch, batch['padding_mask']) + + # Partition padding_mask for context parallel on every PP stage. Partition + # token-like tensors only on stages that own them. + if is_tp_rank_0: + cp_slice_keys = ['padding_mask'] + if is_first_stage or mtp_on_this_rank: + cp_slice_keys.extend(['tokens', 'position_ids']) + if is_last_stage or mtp_on_this_rank: + cp_slice_keys.extend(['labels', 'loss_mask']) + 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. + if is_tp_rank_0: + cu_seqlen_size = torch.tensor(batch['cu_seqlens'].size(0), dtype=torch.int32, device=dev) + else: + cu_seqlen_size = torch.empty(1, dtype=torch.int32, device=dev) + broadcast_tensor(cu_seqlen_size, tp_src_rank, tp_group) + cu_seqlen_size = cu_seqlen_size.item() + + # 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: + total_tokens = torch.tensor(batch['padding_mask'].size(0), dtype=torch.int32, device=dev) + else: + total_tokens = torch.empty(1, dtype=torch.int32, device=dev) + broadcast_tensor(total_tokens, tp_src_rank, tp_group) + total_tokens = total_tokens.item() + + # Step1: Prepare "tokens", "position_ids" on all ranks. + if is_first_stage or mtp_on_this_rank: + if is_tp_rank_0: + assert batch['tokens'].dtype == torch.int64 + assert batch['position_ids'].dtype == torch.int64 + batch['tokens'] = batch['tokens'].view(1, total_tokens) + batch['position_ids'] = batch['position_ids'].view(1, total_tokens) + else: + batch['tokens'] = torch.empty([1, total_tokens], dtype=torch.int64, device=dev) + batch['position_ids'] = torch.empty([1, total_tokens], dtype=torch.int64, device=dev) + else: + # Non first stage rank doesn't need tokens and position_ids. + batch['tokens'] = None + batch['position_ids'] = None + + # Step2: Prepare "labels", "loss_mask" on all ranks. + if is_last_stage or mtp_on_this_rank: + if is_tp_rank_0: + assert batch['labels'].dtype == torch.int64 + assert batch['loss_mask'].dtype == torch.float32 + batch['labels'] = batch['labels'].view(1, total_tokens) + batch['loss_mask'] = batch['loss_mask'].view(1, total_tokens) + else: + batch['labels'] = torch.empty([1, total_tokens], dtype=torch.int64, device=dev) + batch['loss_mask'] = torch.empty([1, total_tokens], dtype=torch.float32, device=dev) + else: + # Non last stage rank doesn't need labels and loss_mask. + batch['labels'] = None + batch['loss_mask'] = None + + # Step3: Prepare "padding_mask" on all TP ranks. + if is_tp_rank_0: + assert batch['padding_mask'].dtype == torch.bool + batch['padding_mask'] = batch['padding_mask'].view(1, total_tokens) + else: + batch['padding_mask'] = torch.empty([1, total_tokens], dtype=torch.bool, device=dev) + + # Step4: Prepare "cu_seqlens", "cu_seqlens_padded", "max_seqlen" on all ranks. + if is_tp_rank_0: + assert batch['cu_seqlens'].dtype == torch.int32 + assert batch['cu_seqlens_padded'].dtype == torch.int32 + assert batch['cu_seqlens'].dim() == 1 + assert batch['cu_seqlens_padded'].dim() == 1 + if type(batch['max_seqlen']) == int: + batch['max_seqlen'] = torch.tensor(batch['max_seqlen'], dtype=torch.int32, device=dev) + else: + assert batch['max_seqlen'].dtype == torch.int32 + assert batch['max_seqlen'].numel() == 1 + else: + batch['cu_seqlens'] = torch.empty([cu_seqlen_size], dtype=torch.int32, device=dev) + batch['cu_seqlens_padded'] = torch.empty([cu_seqlen_size], dtype=torch.int32, device=dev) + batch['max_seqlen'] = torch.empty(1, dtype=torch.int32, device=dev) + + # Broadcast batch inside TP group. + broadcast_tensor(batch['tokens'], tp_src_rank, tp_group) + broadcast_tensor(batch['position_ids'], tp_src_rank, tp_group) + broadcast_tensor(batch['labels'], tp_src_rank, tp_group) + broadcast_tensor(batch['loss_mask'], tp_src_rank, tp_group) + broadcast_tensor(batch['padding_mask'], tp_src_rank, tp_group) + broadcast_tensor(batch['cu_seqlens'], tp_src_rank, tp_group) + broadcast_tensor(batch['cu_seqlens_padded'], tp_src_rank, tp_group) + broadcast_tensor(batch['max_seqlen'], tp_src_rank, tp_group) + + # Extract the data from batch after broadcasting. + tokens = batch['tokens'] + position_ids = batch['position_ids'] + labels = batch['labels'] + loss_mask = batch['loss_mask'] + padding_mask = batch['padding_mask'] + cu_seqlens = batch['cu_seqlens'] + cu_seqlens_padded = batch['cu_seqlens_padded'] + max_seqlen = batch['max_seqlen'].item() + + # Keep original boundaries for loss paths that must identify padding rows. + # Attention kernels and THD partitioning consume the padded boundaries. + packed_seq_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=max_seqlen, + max_seqlen_kv=max_seqlen, + local_cp_size=None, + cp_group=None, + cp_partition_mode=cp_partition_mode, + pad_between_seqs=False, + ) + + if pad_alignment is not None: + tokens, labels, loss_mask, position_ids, packed_seq_params, padding_mask = ( + pad_sequence_for_thd( + tokens, + labels, + loss_mask, + position_ids, + packed_seq_params, + alignment=alignment, + target_len=target_len, + max_num_seqs=max_num_seqs, + pad_by_appending_dummy_seq=getattr( + config, 'pad_packed_seq_by_appending_dummy_seq', True + ), + padding_mask=padding_mask, + cp_group=cp_group, + ) + ) + + # "attention_mask" is not valid for sequence packing, so set it to None. + return tokens, labels, loss_mask, None, position_ids, packed_seq_params, padding_mask diff --git a/megatron/core/datasets/data_schedule_utils.py b/megatron/core/datasets/data_schedule_utils.py new file mode 100644 index 00000000000..137c08511c3 --- /dev/null +++ b/megatron/core/datasets/data_schedule_utils.py @@ -0,0 +1,469 @@ +# Copyright (c) 2025 NVIDIA CORPORATION. All rights reserved. + +from typing import Dict, List, Literal, Optional, Sequence + +import torch + +from megatron.core.extensions.transformer_engine import get_thd_partitioned_indices +from megatron.core.rerun_state_machine import RerunDataIterator + + +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. + + ``zigzag`` uses TE's THD partitioned indices. ``contiguous`` splits the + flattened rows into equal rank-contiguous slices. + + 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: + return + cp_rank = cp_group.rank() + # Partition with padded cumulative lengths so CP slices match the THD + # sequence boundaries consumed by attention kernels. + cu_seqlens = batch["cu_seqlens_padded"] + # 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()) + 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 + 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) + + +def _unpack_batch(batch: List[Dict[str, torch.Tensor]]) -> List[Dict[str, torch.Tensor]]: + """Normalize samples that are already unpacked by the varlen dataset.""" + for sample in batch: + if "padded_seq_len" not in sample: + raise KeyError("sequence packing samples must provide 'padded_seq_len'") + for key, value in sample.items(): + if value.ndim == 2 and value.shape[0] == 1: + sample[key] = value.squeeze(0) + if "original_seq_len" not in sample: + sample["original_seq_len"] = sample["padded_seq_len"].clone() + return batch + + +def _get_global_seqlens_and_ids( + padded_subsample_seqlens: torch.Tensor, original_subsample_seqlens: torch.Tensor, dp_group +): + """ + Gathers the sequence lengths of all subsamples from all DP ranks and calculates global IDs. + """ + # Collect the number of subsamples from all ranks + assert padded_subsample_seqlens.shape == original_subsample_seqlens.shape + num_local_subsamples = padded_subsample_seqlens.shape[0] + local_len = torch.tensor( + [num_local_subsamples], dtype=torch.int32, device=padded_subsample_seqlens.device + ) + dp_subsample_count = [torch.zeros_like(local_len) for _ in range(dp_group.size())] + torch.distributed.all_gather(dp_subsample_count, local_len, group=dp_group) + + # Gather padded and original lengths together so scheduling and FLOPs use + # the same global sample ordering without an extra collective. + dp_subsample_counts = torch.stack(dp_subsample_count, dim=0).cpu().view(-1) + max_sub_samples = int(dp_subsample_counts.max().item()) + local_seqlens = torch.stack([padded_subsample_seqlens, original_subsample_seqlens], dim=1) + + if num_local_subsamples < max_sub_samples: + local_seqlens = torch.cat( + [ + local_seqlens, + torch.zeros( + (max_sub_samples - num_local_subsamples, 2), + dtype=torch.int32, + device=local_seqlens.device, + ), + ], + dim=0, + ) + + seqlens_gathered = [torch.empty_like(local_seqlens) for _ in range(dp_group.size())] + torch.distributed.all_gather(seqlens_gathered, local_seqlens, group=dp_group) + + # Trim each seqlens_gathered to the length of the correct sample + for dp_rank, seqlen in enumerate(seqlens_gathered): + seqlens_gathered[dp_rank] = seqlen[: dp_subsample_counts[dp_rank]] + + seqlens_gathered = torch.cat(seqlens_gathered, dim=0).cpu() + padded_seqlens_gathered = seqlens_gathered[:, 0].tolist() + original_seqlens_gathered = seqlens_gathered[:, 1].tolist() + + # Calculate the offsets to assign unique global ID to each subsample. + csum = torch.cumsum(dp_subsample_counts, dim=0, dtype=torch.int32) + offsets = torch.cat([torch.zeros(1, dtype=torch.int32), csum], dim=0) + + # Calculate global ID for each subsample + dp_rank = dp_group.rank() + global_ids = torch.arange( + len(padded_seqlens_gathered), dtype=torch.int32, device=padded_subsample_seqlens.device + ) + + # Create a list of (global_id, seqlen) tuples for scheduling + global_id_seqlens = [(i, padded_seqlens_gathered[i]) for i in range(len(global_ids))] + + # Get the global IDs locally present on this rank + start_idx = offsets[dp_rank] + end_idx = offsets[dp_rank + 1] + + global_ids_this_rank = global_ids[start_idx:end_idx] + + return ( + global_id_seqlens, + global_ids_this_rank, + offsets, + padded_seqlens_gathered, + original_seqlens_gathered, + ) + + +def _pack_sequences( + samples: List, padded_lengths: torch.Tensor, original_lengths: torch.Tensor, dev: torch.device +) -> Dict[str, torch.Tensor]: + """Pack multiple samples into a single packed sample.""" + + def _pack_tensors(tensors): + return torch.cat([t.reshape(-1) for t in tensors], dim=0) + + new_sample = {} + for key in ['tokens', 'labels', 'loss_mask', 'position_ids']: + if key in samples[0]: + new_sample[key] = _pack_tensors([sample[key] for sample in samples]) + + padded_lengths = padded_lengths.to(device=dev, dtype=torch.int32, non_blocking=True).reshape(-1) + cu_seqlens_padded = torch.empty(padded_lengths.numel() + 1, device=dev, dtype=torch.int32) + cu_seqlens_padded[0] = 0 + cu_seqlens_padded[1:] = torch.cumsum(padded_lengths, dim=0) + max_seqlen = torch.max(padded_lengths).to(dtype=torch.int32) + + new_sample["cu_seqlens_padded"] = cu_seqlens_padded + new_sample["max_seqlen"] = max_seqlen + + original_lengths = original_lengths.to( + device=dev, dtype=torch.int32, non_blocking=True + ).reshape(-1) + cu_seqlens = torch.empty(original_lengths.numel() + 1, device=dev, dtype=torch.int32) + cu_seqlens[0] = 0 + cu_seqlens[1:] = torch.cumsum(original_lengths, dim=0).reshape(-1) + new_sample["cu_seqlens"] = cu_seqlens + + return new_sample + + +def broadcast_tensor(item, src_rank, group) -> None: + """Broadcast a tensor from src_rank to all ranks in the group.""" + if item is not None: + torch.distributed.broadcast(item, src_rank, group=group) + + +def broadcast_scalars(values: List, group, dev, dtype=torch.float32) -> List: + """ + Broadcast scalar values from rank 0 to all ranks in the group. + + Args: + values: List of scalar values to broadcast (only used on rank 0). + group: The process group to broadcast within. + dev: The device to use for the tensor. + dtype: The data type for the tensor. + + Returns: + List of broadcasted values. + """ + if group.size() <= 1: + return values + + src_rank = torch.distributed.get_process_group_ranks(group)[0] + num_values = len(values) + + if group.rank() == 0: + info_to_broadcast = torch.tensor(values, dtype=dtype, device=dev) + else: + info_to_broadcast = torch.zeros(num_values, dtype=dtype, device=dev) + + broadcast_tensor(info_to_broadcast, src_rank, group) + + if group.rank() != 0: + values = info_to_broadcast.cpu().tolist() + + return values + + +def create_data_iterator(new_samples, tp_group, config, vpp_needs_data=None): + """Create independent iterators for the virtual pipeline stages.""" + if ( + config.virtual_pipeline_model_parallel_size is not None + and config.virtual_pipeline_model_parallel_size > 1 + ): + vpp_size = config.virtual_pipeline_model_parallel_size + if tp_group.rank() == 0: + new_data_iterator = [] + for vp_stage in range(vpp_size): + if vpp_needs_data is not None and vpp_needs_data[vp_stage]: + samples = [dict(sample) for sample in new_samples] + new_data_iterator.append(RerunDataIterator(iter(samples))) + else: + metadata_keys = ['max_seqlen', 'cu_seqlens', 'cu_seqlens_padded'] + metadata = [ + {key: sample[key] for key in metadata_keys if key in sample} + for sample in new_samples + ] + new_data_iterator.append(RerunDataIterator(iter(metadata))) + else: + new_data_iterator = [None for _ in range(vpp_size)] + else: + new_data_iterator = RerunDataIterator(iter(new_samples)) if tp_group.rank() == 0 else None + + return new_data_iterator + + +def reroute_samples_to_dcp_ranks( + batch, + global_ids_this_rank, + global_id_seqlens, + sample_id_groups, + offsets, + dp_group, + dp_cp_group, + total_dcp_gpus, +): + """ + Reroutes the sub-samples to the correct rank after scheduling. + + For each key in the batch dict, we perform an all-to-all communication + to transfer the data to the correct ranks. + """ + + dp_global_ranks = torch.distributed.get_process_group_ranks(dp_group) + dp_cp_global_ranks = torch.distributed.get_process_group_ranks(dp_cp_group) + global_to_dcp_rank = {global_rank: rank for rank, global_rank in enumerate(dp_cp_global_ranks)} + + def _gid_to_src_rank(gid: int) -> int: + dp_src_rank = torch.bucketize(gid, offsets[1:] - 1) + return global_to_dcp_rank[dp_global_ranks[int(dp_src_rank)]] + + gid2local_id = {int(gid): i for i, gid in enumerate(global_ids_this_rank)} + dcp_rank = dp_cp_group.rank() + dp_ranks = [global_to_dcp_rank[global_rank] for global_rank in dp_global_ranks] + + data_keys = batch[0].keys() + + # Create the send plan + combined_sample_id_groups: List[List[int]] = [[] for _ in range(total_dcp_gpus)] + for d in range(total_dcp_gpus): + for sample_id_group in sample_id_groups: + combined_sample_id_groups[d].extend(sample_id_group[d]) + for dest_rank in range(total_dcp_gpus): + combined_sample_id_groups[dest_rank].sort() + + send_ids_sorted = [ + gid for d in dp_ranks for gid in combined_sample_id_groups[d] if gid in global_ids_this_rank + ] + + send_num_split = [0] * total_dcp_gpus + send_lens_split = [0] * total_dcp_gpus + for dest_rank in range(total_dcp_gpus): + if dest_rank in dp_ranks: + send_seq_lens = [ + global_id_seqlens[gid][1] + for gid in combined_sample_id_groups[dest_rank] + if gid in global_ids_this_rank + ] + send_num_split[dest_rank] = len(send_seq_lens) + send_lens_split[dest_rank] = sum(send_seq_lens) + else: + send_lens_split[dest_rank] = 0 + + # Create the recv plan + recv_sample_id_groups = [[] for _ in range(total_dcp_gpus)] + for gid in combined_sample_id_groups[dcp_rank]: + src_rank = _gid_to_src_rank(gid) + recv_sample_id_groups[src_rank].append(gid) + + recv_lens_split = [0] * total_dcp_gpus + for src_rank in range(total_dcp_gpus): + recv_lens_split[src_rank] = sum( + [global_id_seqlens[gid][1] for gid in recv_sample_id_groups[src_rank]] + ) + + recv_ids_sorted = [gid for d in range(total_dcp_gpus) for gid in recv_sample_id_groups[d]] + recv_counts = [len(recv_sample_id_groups[d]) for d in range(total_dcp_gpus)] + + recv_samples = [{k: None for k in data_keys} for _ in range(sum(recv_counts))] + + def _pack_sample_by_key(key: str) -> torch.Tensor: + flattened_tensors = [] + for gid in send_ids_sorted: + t = batch[gid2local_id[gid]][key].to(torch.cuda.current_device(), non_blocking=True) + flattened_tensors.append(t.reshape(-1)) + return ( + torch.cat(flattened_tensors, dim=0) + if flattened_tensors + else torch.empty(0, device=torch.cuda.current_device(), dtype=batch[0][key].dtype) + ) + + def _unpack_sample_by_key(key: str, recv_tensor: torch.Tensor): + cursor = 0 + for i, gid in enumerate(recv_ids_sorted): + sample_len = ( + 1 if key in ["original_seq_len", "padded_seq_len"] else global_id_seqlens[gid][1] + ) + recv_samples[i][key] = recv_tensor[cursor : cursor + sample_len] + cursor += sample_len + + for key in data_keys: + output_split_sizes, input_split_sizes = ( + (recv_counts, send_num_split) + if key in ["original_seq_len", "padded_seq_len"] + else (recv_lens_split, send_lens_split) + ) + send_tensor = _pack_sample_by_key(key) + recv_tensor_size = sum(output_split_sizes) + recv_tensor = torch.empty( + recv_tensor_size, device=torch.cuda.current_device(), dtype=send_tensor.dtype + ) + torch.distributed.all_to_all_single( + output=recv_tensor, + input=send_tensor, + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=dp_cp_group, + ) + _unpack_sample_by_key(key, recv_tensor) + + recv_sample_with_id = {recv_id: recv_samples[i] for i, recv_id in enumerate(recv_ids_sorted)} + return recv_sample_with_id + + +def build_packed_microbatches( + grouped_samples: List[List[Dict[str, torch.Tensor]]], dev: torch.device +) -> List[Dict[str, torch.Tensor]]: + """Build packed samples for each microbatch.""" + num_micro_batches = len(grouped_samples) + seg_starts: List[int] = [0] + original_lens_tensors = [] + padded_lens_tensors = [] + + for i in range(num_micro_batches): + samples = grouped_samples[i] + seg_starts.append(seg_starts[-1] + len(samples)) + original_lens_tensors.extend([s["original_seq_len"].reshape(-1) for s in samples]) + padded_lens_tensors.extend([s["padded_seq_len"].reshape(-1) for s in samples]) + + padded_lens_all_gpu = torch.cat(padded_lens_tensors, dim=0).to(dtype=torch.int32) + original_lens_all_gpu = torch.cat(original_lens_tensors, dim=0).to(dtype=torch.int32) + + new_samples: List[Dict[str, torch.Tensor]] = [] + for i in range(num_micro_batches): + samples = grouped_samples[i] + lens_padded = padded_lens_all_gpu[seg_starts[i] : seg_starts[i + 1]] + lens_original = original_lens_all_gpu[seg_starts[i] : seg_starts[i + 1]] + new_sample = _pack_sequences(samples, lens_padded, lens_original, dev) + new_samples.append(new_sample) + + return new_samples + + +def get_batch_and_global_seqlens(data_iterator, num_microbatches, dp_group): + """ + Get the batch and global sequence lengths. + Each DP rank loads the same number of sequences, so we need to gather the sequence + lengths from all ranks then we can schedule the sequences into groups. + Args: + data_iterator: The data iterator. + num_microbatches: The number of microbatches. + dp_group: The data parallel group. + + Returns: + batch: The batch. + global_id_seqlens: The global sequence lengths. + global_ids_this_rank: The global IDs locally present on this rank. + """ + + batch_list = [next(data_iterator) for _ in range(num_microbatches)] + + batch = [] + for item in batch_list: + if isinstance(item, dict): + batch.append(item) + elif isinstance(item, list): + batch.extend(item) + else: + raise ValueError(f"Invalid item type: {type(item)}") + + # Normalize the optional leading batch dimension before scheduling. + batch = _unpack_batch(batch) + + padded_subsample_seqlens = torch.cat([sample["padded_seq_len"] for sample in batch]).to( + dtype=torch.int32, device=torch.cuda.current_device() + ) + original_subsample_seqlens = torch.cat([sample["original_seq_len"] for sample in batch]).to( + dtype=torch.int32, device=torch.cuda.current_device() + ) + + ( + global_id_seqlens, + global_ids_this_rank, + offsets, + padded_seqlens_gathered, + original_seqlens_gathered, + ) = _get_global_seqlens_and_ids(padded_subsample_seqlens, original_subsample_seqlens, dp_group) + + return ( + batch, + global_id_seqlens, + global_ids_this_rank, + offsets, + padded_seqlens_gathered, + original_seqlens_gathered, + ) diff --git a/megatron/core/datasets/readme.md b/megatron/core/datasets/readme.md index 58721b7471b..40cdc111329 100644 --- a/megatron/core/datasets/readme.md +++ b/megatron/core/datasets/readme.md @@ -204,6 +204,26 @@ If the later training job does not specify `--global-batch-size` (which is neede `tools/prepare_cache.py` does not support `--mock-data`, `--sft`, `--fim-data`, or `--step-batch-size-schedule`. +## Packing Scheduler + +The packing scheduler reschedules variable-length sequences across DPxCP ranks to improve GPU utilization. It is built around the following modules: + +### `data_schedule` + +This module contains the high-level scheduling logic and entry points: + +- **`BasePackingScheduler`**: Abstract base class for packing schedulers. Defines the interface for `get_groups_and_subsamples()` (scheduling algorithm) and `run()` (full scheduling pipeline including fetch, schedule, reroute, pack, TP synchronization, and VPP handling). + +- **`DpBalancedScheduler`**: A concrete scheduler that packs sequences in their original order until reaching the max sequence length limit per DPxCP rank. Supports aligning the number of microbatches to DP size and VPP stage multiples. + +- **`wrap_data_iterator()`**: Top-level entry point that wraps an existing `data_iterator`. Every TP-rank-0 PP stage schedules its local iterator, while scalar schedule results are synchronized inside each TP group. It returns the packed iterator, updated number of microbatches, and FLOPs statistics. + +- **`get_batch_on_this_rank_for_sequence_packing()`**: Fetches a packed microbatch on TP rank 0, broadcasts it within the TP group, constructs `PackedSeqParams` (with `cu_seqlens`, `max_seqlen`, `qkv_format=thd`), and optionally partitions sequences across CP ranks using Transformer Engine's `thd_get_partitioned_indices`. + +### `data_schedule_utils.py` + +This module contains the utility functions used by the schedulers. + ## Fast DataLoader initialization Especially for large-scale runs, DataLoader initialization can take several minutes, since it involves opening and memory-mapping multiple files and can significantly stress the filesystem. To speed up this process, we have developed the following three optimizations, controlled by configuration flags: diff --git a/megatron/core/distributed/finalize_model_grads.py b/megatron/core/distributed/finalize_model_grads.py index 12253e6c4b6..fb178b86893 100644 --- a/megatron/core/distributed/finalize_model_grads.py +++ b/megatron/core/distributed/finalize_model_grads.py @@ -319,7 +319,11 @@ def reset_model_temporary_tensors(config: TransformerConfig, model: List[torch.n """ for model_chunk in model: for module in get_attr_wrapped_model(model_chunk, 'modules')(): - if config.moe_router_enable_expert_bias and hasattr(module, 'expert_bias'): + if ( + config.moe_router_enable_expert_bias + and hasattr(module, 'expert_bias') + and module.expert_bias is not None + ): module.local_tokens_per_expert.zero_() if ( config.moe_router_load_balancing_type == "global_aux_loss" @@ -351,6 +355,7 @@ def _update_router_expert_bias( if ( hasattr(module, 'expert_bias') and module.training + and module.expert_bias is not None and not getattr(module, 'frozen_expert_bias', False) ): tokens_per_expert_list.append(module.local_tokens_per_expert) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 3a8d7f23d09..f622e160f03 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -1748,11 +1748,11 @@ 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. - # tokens_per_sample is only for MoE sequence-level aux loss reshaping. + # 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("tokens_per_sample") + self.kept_packed_seq_params.discard("cp_partition_mode") if get_te_version() < PkgVersion("2.2.0"): self.kept_packed_seq_params.discard("pad_between_seqs") @@ -2006,12 +2006,23 @@ def __init__( tp_group_for_te = None if is_te_min_version("2.14.0"): - extra_kwargs["single_grouped_weight"] = getattr( - config, "moe_single_grouped_weight", False - ) - extra_kwargs["single_grouped_bias"] = getattr( - config, "moe_single_grouped_bias", False - ) + # Some TE 2.14 builds predate these keyword arguments. Cache the + # installed constructor signature instead of relying on the version alone. + global _TE_GROUPED_LINEAR_INIT_PARAMS + try: + grouped_linear_init_params = _TE_GROUPED_LINEAR_INIT_PARAMS + except NameError: + grouped_linear_init_params = _TE_GROUPED_LINEAR_INIT_PARAMS = set( + inspect.signature(te.pytorch.GroupedLinear.__init__).parameters + ) + if "single_grouped_weight" in grouped_linear_init_params: + extra_kwargs["single_grouped_weight"] = getattr( + config, "moe_single_grouped_weight", False + ) + if "single_grouped_bias" in grouped_linear_init_params: + extra_kwargs["single_grouped_bias"] = getattr( + config, "moe_single_grouped_bias", False + ) self.te_quant_params: Optional[TEQuantizationParams] = None quant_config = get_quant_config_or_none(name, config.quant_recipe) @@ -3396,3 +3407,16 @@ def set_save_original_input(module): from transformer_engine.pytorch.float8_tensor import Float8Tensor except ImportError: Float8Tensor = None + + +def get_thd_partitioned_indices( + cu_seqlens: torch.Tensor, total_tokens: int, cp_size: int, cp_rank: int +) -> torch.Tensor: + """Get partitioned indices for THD data in context parallelism.""" + assert is_te_min_version("1.10.0"), ( + "Please update Transformer Engine to >= 1.10 to use " + "Context Parallel with THD format data" + ) + import transformer_engine_torch as tex + + return tex.thd_get_partitioned_indices(cu_seqlens, total_tokens, cp_size, cp_rank) diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py index 5411b676d83..e0eb53f7506 100644 --- a/megatron/core/fp8_utils.py +++ b/megatron/core/fp8_utils.py @@ -831,6 +831,29 @@ def get_fp8_context(config: TransformerConfig, layer_no: int = -1, is_init: bool return fp8_context + def get_fp8_disabled_context(config: TransformerConfig, is_init: bool = False): + """Return a context manager that disables TE quantization. + + Use this around submodule construction or execution that must stay in a higher + precision while its enclosing module uses an FP8 or FP4 context. + + Args: + config: Transformer configuration that controls quantization. + is_init: Whether to disable the parameter-initialization context instead of + the forward autocast context. + + Returns: + A disabled TE quantization context when quantization is active, otherwise a + no-op context. + """ + if is_init: + if not (config.fp8_param or config.fp4_param): + return nullcontext() + return transformer_engine.pytorch.fp8_model_init(enabled=False) + if not (config.fp8 or config.fp4): + return nullcontext() + return transformer_engine.pytorch.fp8_autocast(enabled=False) + else: def get_fp8_recipe(config: TransformerConfig): @@ -841,6 +864,10 @@ def get_fp8_context(config: TransformerConfig, layer_no: int = -1, is_init: bool """Returns dummy fp8 context manager since TE is not available.""" return nullcontext() + def get_fp8_disabled_context(config: TransformerConfig, is_init: bool = False): + """Return a no-op context manager since TE is not available.""" + return nullcontext() + if HAVE_TE: from transformer_engine.pytorch.fp8 import FP8GlobalStateManager diff --git a/megatron/core/fusions/fused_bias_dropout.py b/megatron/core/fusions/fused_bias_dropout.py index 2eb4007f75c..92466ba7391 100644 --- a/megatron/core/fusions/fused_bias_dropout.py +++ b/megatron/core/fusions/fused_bias_dropout.py @@ -1,10 +1,13 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. -from typing import Optional, Tuple +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from typing import TYPE_CHECKING, Optional, Tuple import torch from megatron.core.jit import jit_fuser +if TYPE_CHECKING: + from megatron.core.tensor_parallel.random import CheckpointWithoutOutputManager + # pylint: disable=missing-function-docstring @@ -80,7 +83,26 @@ def bias_dropout_add_fused_inference( return _bias_dropout_add_func(x_with_bias, residual, prob, False) -def get_bias_dropout_add(training, fused): +def get_bias_dropout_add( + training, fused, mhc_recompute_manager: Optional['CheckpointWithoutOutputManager'] = None +): + """ + Get the bias-dropout-add function. + + Args: + training: Whether in training mode. + fused: Whether to use fused implementation. + mhc_recompute_manager: Optional CheckpointWithoutOutputManager for checkpoint management. + When provided, the returned function will wrap the BDA operation with + CheckpointWithoutOutput for memory-efficient recomputation. + + Returns: + A callable that performs bias-dropout-add operation. + """ + if mhc_recompute_manager is not None: + # Return a checkpointed version that handles tuple unpacking internally + return _get_checkpointed_bda(training, fused, mhc_recompute_manager) + if fused: # jit scripting for a nn.module (with dropout) is not # triggering the fusion kernel. For now, we use two @@ -92,3 +114,68 @@ def get_bias_dropout_add(training, fused): return bias_dropout_add_fused_inference else: return bias_dropout_add_unfused(training) + + +def _get_checkpointed_bda(training, fused, mhc_recompute_manager: 'CheckpointWithoutOutputManager'): + """ + Create a checkpointed bias-dropout-add function. + + This function handles: + 1. Tuple unpacking for x_with_bias (required because save_for_backward can't save tuples) + 2. Non-tensor arguments like dropout probability (handled by CheckpointWithoutOutput) + 3. Auto-registration to the CheckpointWithoutOutputManager + + Args: + training: Whether in training mode. + fused: Whether to use fused implementation. + mhc_recompute_manager: CheckpointWithoutOutputManager for checkpoint management. + + Returns: + A callable that performs checkpointed bias-dropout-add operation. + """ + from megatron.core.tensor_parallel.random import CheckpointWithoutOutput + + # Get the underlying BDA function + if fused: + if training: + bda_func = bias_dropout_add_fused_train + else: + bda_func = bias_dropout_add_fused_inference + else: + bda_func = bias_dropout_add_unfused(training) + + def _checkpointed_bda(x_with_bias, residual, prob): + """ + Checkpointed BDA that handles tuple unpacking internally. + + Args: + x_with_bias: Either a tuple (x, bias) or a single tensor x. + residual: Residual tensor. + prob: Dropout probability. + + Returns: + Output tensor after bias-dropout-add. + """ + # Create checkpoint with manager + ckpt = CheckpointWithoutOutput(ckpt_manager=mhc_recompute_manager) + + # Handle case where x_with_bias might be a single tensor (e.g., from IdentityOp) + if isinstance(x_with_bias, tuple): + x, bias = x_with_bias + else: + x = x_with_bias + bias = None + + # Wrapper function that re-packs the tuple for the actual BDA function + def _bda_wrapper(output, bias, res, dropout): + return bda_func((output, bias), res, dropout) + + # Call checkpoint with unpacked arguments + result = ckpt.checkpoint(_bda_wrapper, x, bias, residual, prob) + + # No-op when manager is set - manager handles all discarding uniformly + ckpt.discard_output_and_register_recompute(result) + + return result + + return _checkpointed_bda diff --git a/megatron/core/fusions/fused_bias_swiglu.py b/megatron/core/fusions/fused_bias_swiglu.py index 632470876c9..8d64e3016a9 100644 --- a/megatron/core/fusions/fused_bias_swiglu.py +++ b/megatron/core/fusions/fused_bias_swiglu.py @@ -48,6 +48,31 @@ def weighted_swiglu(y, weights): return res.to(dtype) +@jit_fuser +def clamped_swiglu(y, clamp_value): + """Perform SwiGLU after clamping both halves of the input.""" + dtype = y.dtype + y_1, y_2 = torch.chunk(y.to(torch.float32), 2, -1) + y_1 = y_1.clamp(min=None, max=clamp_value) + y_2 = y_2.clamp(min=-clamp_value, max=clamp_value) + res = F.silu(y_1) * y_2 + return res.to(dtype) + + +@jit_fuser +def bias_clamped_swiglu(y, bias, clamp_value): + """Perform clamped SwiGLU after bias addition.""" + return clamped_swiglu(y + bias, clamp_value) + + +@jit_fuser +def clamped_weighted_swiglu(y, weights, clamp_value): + """Perform token-weighted clamped SwiGLU.""" + dtype = y.dtype + res = clamped_swiglu(y, clamp_value) * weights + return res.to(dtype) + + # gradient of tanh approximation of gelu # gradient of actual gelu is: # 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) @@ -97,12 +122,50 @@ def weighted_swiglu_back(g, y, weights): return input_grad.to(input_dtype), weights_grad.to(w_dtype) +@jit_fuser +def clamped_swiglu_back(g, y, clamp_value): + """Compute the input gradient for clamped SwiGLU.""" + dtype = y.dtype + y_1, y_2 = torch.chunk(y.to(torch.float32), 2, -1) + y_1c = y_1.clamp(min=None, max=clamp_value) + y_2c = y_2.clamp(min=-clamp_value, max=clamp_value) + res = torch.cat( + ( + g + * torch.sigmoid(y_1c) + * (1 + y_1c * (1 - torch.sigmoid(y_1c))) + * y_2c + * (y_1 <= clamp_value).to(g.dtype), + g * F.silu(y_1c) * ((y_2 >= -clamp_value) & (y_2 <= clamp_value)).to(g.dtype), + ), + -1, + ) + return res.to(dtype) + + +@jit_fuser +def bias_clamped_swiglu_back(g, y, bias, clamp_value): + """Compute the input gradient for clamped SwiGLU with bias.""" + return clamped_swiglu_back(g, y + bias, clamp_value) + + +@jit_fuser +def clamped_weighted_swiglu_back(g, y, weights, clamp_value): + """Compute input and weight gradients for token-weighted clamped SwiGLU.""" + input_dtype = y.dtype + w_dtype = weights.dtype + input_grad = clamped_swiglu_back(g * weights, y, clamp_value) + weights_grad = clamped_swiglu(y, clamp_value) * g.to(w_dtype) + weights_grad = torch.sum(weights_grad, dim=-1, keepdim=True) + return input_grad.to(input_dtype), weights_grad.to(w_dtype) + + class BiasSwiGLUFunction(torch.autograd.Function): """Custom autograd function for SwiGLU activation with bias support.""" @staticmethod @nvtx_decorator() - def forward(ctx, input, bias, fp8_input_store, cpu_offload_input): + def forward(ctx, input, bias, fp8_input_store, cpu_offload_input, clamp_value): """Forward pass of biased SwiGLU activation. Args: @@ -121,6 +184,9 @@ def forward(ctx, input, bias, fp8_input_store, cpu_offload_input): ctx.save_for_backward(input_for_backward, bias) ctx.ori_input_dtype = input.dtype ctx.fp8_input_store = fp8_input_store + ctx.clamp_value = clamp_value + if clamp_value is not None and clamp_value > 0: + return bias_clamped_swiglu(input, bias, clamp_value) return bias_swiglu(input, bias) @staticmethod @@ -140,8 +206,11 @@ def backward(ctx, grad_output): """ input, bias = ctx.saved_tensors input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input - tmp = bias_swiglu_back(grad_output, input, bias) - return tmp, tmp, None, None + if ctx.clamp_value is not None and ctx.clamp_value > 0: + tmp = bias_clamped_swiglu_back(grad_output, input, bias, ctx.clamp_value) + else: + tmp = bias_swiglu_back(grad_output, input, bias) + return tmp, tmp, None, None, None class SwiGLUFunction(torch.autograd.Function): @@ -149,7 +218,7 @@ class SwiGLUFunction(torch.autograd.Function): @staticmethod @nvtx_decorator() - def forward(ctx, input, fp8_input_store, cpu_offload_input): + def forward(ctx, input, fp8_input_store, cpu_offload_input, clamp_value): """Forward pass of SwiGLU activation. Args: @@ -166,6 +235,9 @@ def forward(ctx, input, fp8_input_store, cpu_offload_input): ctx.save_for_backward(input_for_backward) ctx.ori_input_dtype = input.dtype ctx.fp8_input_store = fp8_input_store + ctx.clamp_value = clamp_value + if clamp_value is not None and clamp_value > 0: + return clamped_swiglu(input, clamp_value) return swiglu(input) @staticmethod @@ -184,29 +256,37 @@ def backward(ctx, grad_output): """ input = ctx.saved_tensors[0] input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input - tmp = swiglu_back(grad_output, input) - return tmp, None, None + if ctx.clamp_value is not None and ctx.clamp_value > 0: + tmp = clamped_swiglu_back(grad_output, input, ctx.clamp_value) + else: + tmp = swiglu_back(grad_output, input) + return tmp, None, None, None class WeightedSwiGLUFunction(torch.autograd.Function): @staticmethod - # bias is an optional argument - def forward(ctx, input, weights, fp8_input_store): + def forward(ctx, input, weights, fp8_input_store, clamp_value): input_for_backward = input.to(torch.float8_e4m3fn) if fp8_input_store else input ctx.save_for_backward(input_for_backward, weights) ctx.ori_input_dtype = input.dtype ctx.fp8_input_store = fp8_input_store + ctx.clamp_value = clamp_value + if clamp_value is not None and clamp_value > 0: + return clamped_weighted_swiglu(input, weights, clamp_value) return weighted_swiglu(input, weights) @staticmethod def backward(ctx, grad_output): input, weights = ctx.saved_tensors input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input - tmp, wgrad = weighted_swiglu_back(grad_output, input, weights) - return tmp, wgrad, None + if ctx.clamp_value is not None and ctx.clamp_value > 0: + tmp, wgrad = clamped_weighted_swiglu_back(grad_output, input, weights, ctx.clamp_value) + else: + tmp, wgrad = weighted_swiglu_back(grad_output, input, weights) + return tmp, wgrad, None, None -def bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False): +def bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False, clamp_value=None): """Implementation of biased SwiGLU that handles different input shapes. This function reshapes the input if necessary, applies the SwiGLU activation @@ -218,6 +298,10 @@ def bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False uses the bias-free SwiGLU variant. fp8_input_store (bool, optional): Whether to store intermediate values in FP8 format. Defaults to False. + cpu_offload_input (bool, optional): Whether to mark saved activation inputs for CPU + offloading. Defaults to False. + clamp_value (float, optional): Maximum gate value and absolute linear value. When None, + preserve the legacy unclamped SwiGLU behavior. Returns: torch.Tensor: Result of biased SwiGLU activation. @@ -229,14 +313,16 @@ def bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False assert len(ori_shape) in [2, 3] input = input.view(-1, ori_shape[-1]) if bias is not None: - output = BiasSwiGLUFunction.apply(input, bias, fp8_input_store, cpu_offload_input) + output = BiasSwiGLUFunction.apply( + input, bias, fp8_input_store, cpu_offload_input, clamp_value + ) else: - output = SwiGLUFunction.apply(input, fp8_input_store, cpu_offload_input) + output = SwiGLUFunction.apply(input, fp8_input_store, cpu_offload_input, clamp_value) return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1) -def weighted_bias_swiglu_impl(input, bias, weights, fp8_input_store=False): +def weighted_bias_swiglu_impl(input, bias, weights, fp8_input_store=False, clamp_value=None): """ Token-wise-weighted bias swiglu fusion. """ @@ -246,7 +332,7 @@ def weighted_bias_swiglu_impl(input, bias, weights, fp8_input_store=False): if bias is not None: raise NotImplementedError("Bias is not supported for weighted swiglu fusion") else: - output = WeightedSwiGLUFunction.apply(input, weights, fp8_input_store) + output = WeightedSwiGLUFunction.apply(input, weights, fp8_input_store, clamp_value) return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py new file mode 100644 index 00000000000..92c436690a6 --- /dev/null +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -0,0 +1,3397 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Fused kernels for mHC (Manifold-Constrained Hyper-Connections). + +Uses Triton and cuda.tile (cuTile) kernels when available, with PyTorch +reference implementations as fallback. Reference (non-fused) implementations +live in ``megatron.core.transformer.hyper_connection`` and are used when fused +kernels are unavailable or when the ``use_fused_mhc`` config flag is False. + +Four fused operations: + - sinkhorn: Sinkhorn-Knopp projection to doubly stochastic matrix + - h_aggregate: weighted n-stream -> 1-stream aggregation + - h_post_bda: fused H_res.T @ residual + H_post * (x + bias) + - proj_rms: fused projection + RMS normalization +""" + +import logging +import math +import os +import shutil +import subprocess +import warnings +from typing import Optional, Tuple + +import torch +from torch import Tensor + +from megatron.core._rank_utils import log_single_rank, safe_get_rank + +logger = logging.getLogger(__name__) +LOG2E = math.log2(math.e) + + +def _env_flag(name: str) -> bool: + return os.getenv(name, "0").lower() in ("1", "true", "yes", "on") + + +def _forced_backend() -> Tuple[str, Optional[Exception]]: + value = os.getenv("MHC_FORCE_BACKEND", "auto").strip().lower() + value = value.replace("-", "_").replace("+", "_") + aliases = { + "auto": "auto", + "mixed": "auto", + "default": "auto", + "native": "native", + "torch": "native", + "pytorch": "native", + "none": "native", + "triton": "triton", + "triton_native": "triton", + "cutile": "cutile", + "cu_tile": "cutile", + "cuda_tile": "cutile", + } + if value not in aliases: + valid = ", ".join(sorted(aliases)) + return "auto", ValueError( + f"Unsupported MHC_FORCE_BACKEND={value!r}; expected one of: {valid}" + ) + return aliases[value], None + + +# --------------------------------------------------------------------------- +# Check cuTile availability +# --------------------------------------------------------------------------- +_CUTILE_AVAILABLE = False +_CUTILE_EXPERIMENTAL_AVAILABLE = False +_CUTILE_DEVICE_SUPPORT_CACHE: Optional[bool] = None +_CUTILE_DEVICE_SUPPORT_ERROR: Optional[str] = None +try: + import cuda.tile as ct + + _CUTILE_AVAILABLE = True + try: + import cuda.tile_experimental as ct_experimental + + _CUTILE_EXPERIMENTAL_AVAILABLE = True + except ImportError: + pass +except ImportError: + pass + + +# --------------------------------------------------------------------------- +# Check Triton availability +# --------------------------------------------------------------------------- +_TRITON_AVAILABLE = False +try: + import triton + import triton.language as tl + + _TRITON_AVAILABLE = True +except ImportError: + pass + + +_MHC_FORCED_BACKEND, _MHC_BACKEND_VALIDATION_ERROR = _forced_backend() + + +def _record_mhc_backend_validation_error(error: Exception) -> None: + global _MHC_BACKEND_VALIDATION_ERROR + if _MHC_BACKEND_VALIDATION_ERROR is None: + _MHC_BACKEND_VALIDATION_ERROR = error + + +def _raise_mhc_backend_validation_error() -> None: + if _MHC_BACKEND_VALIDATION_ERROR is not None: + raise _MHC_BACKEND_VALIDATION_ERROR + if _MHC_FORCED_BACKEND == "cutile" and not is_cutile_available(): + raise RuntimeError( + "MHC_FORCE_BACKEND=cutile was requested, but cuTile does not support " + f"the current device: {_CUTILE_DEVICE_SUPPORT_ERROR}" + ) + + +if _MHC_FORCED_BACKEND == "native": + _TRITON_AVAILABLE = False + _CUTILE_AVAILABLE = False + _CUTILE_EXPERIMENTAL_AVAILABLE = False +elif _MHC_FORCED_BACKEND == "triton": + if not _TRITON_AVAILABLE: + _record_mhc_backend_validation_error( + RuntimeError("MHC_FORCE_BACKEND=triton was requested, but Triton is not available") + ) + _CUTILE_AVAILABLE = False + _CUTILE_EXPERIMENTAL_AVAILABLE = False +elif _MHC_FORCED_BACKEND == "cutile": + if not _CUTILE_AVAILABLE: + _record_mhc_backend_validation_error( + RuntimeError("MHC_FORCE_BACKEND=cutile was requested, but cuTile is not available") + ) + _TRITON_AVAILABLE = False + +if _env_flag("MHC_DISABLE_TRITON"): + if _MHC_FORCED_BACKEND == "triton": + _record_mhc_backend_validation_error( + ValueError("MHC_FORCE_BACKEND=triton conflicts with MHC_DISABLE_TRITON=1") + ) + _TRITON_AVAILABLE = False + +if _env_flag("MHC_DISABLE_CUTILE"): + if _MHC_FORCED_BACKEND == "cutile": + _record_mhc_backend_validation_error( + ValueError("MHC_FORCE_BACKEND=cutile conflicts with MHC_DISABLE_CUTILE=1") + ) + _CUTILE_AVAILABLE = False + _CUTILE_EXPERIMENTAL_AVAILABLE = False + + +def is_cutile_available() -> bool: + """Return True if cuTile fused kernels are enabled.""" + return _CUTILE_AVAILABLE and _cutile_supports_current_device() + + +def _get_tileiras_path() -> Optional[str]: + """Return the tileiras compiler path if it can be found.""" + tileiras = shutil.which("tileiras") + if tileiras is not None: + return tileiras + + cuda_home = os.getenv("CUDA_HOME") or os.getenv("CUDA_PATH") or "/usr/local/cuda" + candidate = os.path.join(cuda_home, "bin", "tileiras") + if os.path.exists(candidate): + return candidate + return None + + +def _cutile_supports_current_device() -> bool: + """Return whether cuTile can compile for the current CUDA device.""" + global _CUTILE_DEVICE_SUPPORT_CACHE, _CUTILE_DEVICE_SUPPORT_ERROR + + if not _CUTILE_AVAILABLE: + return False + if _CUTILE_DEVICE_SUPPORT_CACHE is not None: + return _CUTILE_DEVICE_SUPPORT_CACHE + + if not torch.cuda.is_available(): + _CUTILE_DEVICE_SUPPORT_ERROR = "CUDA is not available" + _CUTILE_DEVICE_SUPPORT_CACHE = False + return False + + major, minor = torch.cuda.get_device_capability() + arch = f"sm_{major}{minor}" + tileiras = _get_tileiras_path() + if tileiras is None: + _CUTILE_DEVICE_SUPPORT_ERROR = "tileiras compiler was not found" + _CUTILE_DEVICE_SUPPORT_CACHE = False + return False + + try: + result = subprocess.run( + [tileiras, "--gpu-name", arch], capture_output=True, check=False, text=True, timeout=10 + ) + except (OSError, subprocess.SubprocessError) as exc: + _CUTILE_DEVICE_SUPPORT_ERROR = str(exc) + _CUTILE_DEVICE_SUPPORT_CACHE = False + return False + + output = f"{result.stdout}\n{result.stderr}" + if "Cannot find option named" in output and arch in output: + _CUTILE_DEVICE_SUPPORT_ERROR = output.strip().splitlines()[0] + _CUTILE_DEVICE_SUPPORT_CACHE = False + return False + + _CUTILE_DEVICE_SUPPORT_CACHE = True + return True + + +def is_triton_available() -> bool: + """Return True if Triton is enabled for supported mHC kernels.""" + return _TRITON_AVAILABLE + + +# ============================================================================ +# Triton implementations (only defined when triton is available) +# ============================================================================ + +if _TRITON_AVAILABLE: + TLOG2E = tl.constexpr(LOG2E) + + # ============================================================================ + # Sinkhorn-Knopp + # ============================================================================ + + @triton.autotune( + configs=[triton.Config({}, num_warps=nw) for nw in (1, 2, 4, 8)], key=["HC", "NUM_ITERS"] + ) + @triton.jit + def _triton_sinkhorn_fwd_kernel( + inp_ptr, out_ptr, M_init_ptr, N_batch, eps, HC: tl.constexpr, NUM_ITERS: tl.constexpr + ): + """Grid: (N_batch,). Each program handles one [HC, HC] matrix.""" + pid = tl.program_id(0) + if pid >= N_batch: + return + + base = pid * HC * HC + offs_r = tl.arange(0, HC) + offs_c = tl.arange(0, HC) + mat_ptrs = base + offs_r[:, None] * HC + offs_c[None, :] + + logits = tl.load(inp_ptr + mat_ptrs).to(tl.float32) + row_max = tl.max(logits, axis=1) + # Subtract row_max before exp to keep the exponent numerically stable. + M = tl.exp2((logits - row_max[:, None]) * TLOG2E) + tl.store(M_init_ptr + mat_ptrs, M.to(M_init_ptr.dtype.element_ty)) + + row_sum = tl.sum(M, axis=1) + M = M / row_sum[:, None] + eps + col_sum = tl.sum(M, axis=0) + M = M / (col_sum[None, :] + eps) + for _ in range(NUM_ITERS - 1): + row_sum = tl.sum(M, axis=1) + M = M / (row_sum[:, None] + eps) + col_sum = tl.sum(M, axis=0) + M = M / (col_sum[None, :] + eps) + + tl.store(out_ptr + mat_ptrs, M.to(out_ptr.dtype.element_ty)) + + @triton.autotune( + configs=[triton.Config({}, num_warps=nw) for nw in (1, 2, 4, 8)], key=["HC", "NUM_ITERS"] + ) + @triton.jit + def _triton_sinkhorn_bwd_kernel( + grad_out_ptr, + M_init_ptr, + grad_inp_ptr, + ws_M_ptr, + ws_rs_ptr, + ws_cs_ptr, + N_batch, + eps, + HC: tl.constexpr, + NUM_ITERS: tl.constexpr, + ): + """Grid: (N_batch,). Each program handles one [HC, HC] backward.""" + pid = tl.program_id(0) + if pid >= N_batch: + return + + base = pid * HC * HC + M_ws_base = pid * 2 * NUM_ITERS * HC * HC + v_ws_base = pid * NUM_ITERS + offs_r = tl.arange(0, HC) + offs_c = tl.arange(0, HC) + mat_ptrs = base + offs_r[:, None] * HC + offs_c[None, :] + + M = tl.load(M_init_ptr + mat_ptrs).to(tl.float32) + for t in range(NUM_ITERS): + ws_off = M_ws_base + (2 * t) * HC * HC + tl.store(ws_M_ptr + ws_off + offs_r[:, None] * HC + offs_c[None, :], M) + + row_sum = tl.sum(M, axis=1) + tl.store(ws_rs_ptr + (v_ws_base + t) * HC + offs_r, row_sum) + if t == 0: + M = M / row_sum[:, None] + eps + else: + M = M / (row_sum[:, None] + eps) + + ws_off = M_ws_base + (2 * t + 1) * HC * HC + tl.store(ws_M_ptr + ws_off + offs_r[:, None] * HC + offs_c[None, :], M) + + col_sum = tl.sum(M, axis=0) + tl.store(ws_cs_ptr + (v_ws_base + t) * HC + offs_c, col_sum) + M = M / (col_sum[None, :] + eps) + + # M is the final forward output. It is the right value for the first VJP + # through the last column-normalization step. + grad = tl.load(grad_out_ptr + mat_ptrs).to(tl.float32) + for t_rev in range(NUM_ITERS): + t = NUM_ITERS - 1 - t_rev + + col_s = tl.load(ws_cs_ptr + (v_ws_base + t) * HC + offs_c).to(tl.float32) + grad = grad / (col_s[None, :] + eps) + col_corr = tl.sum(grad * M, axis=0) + grad = grad - col_corr[None, :] + M = tl.load( + ws_M_ptr + + M_ws_base + + (2 * t + 1) * HC * HC + + offs_r[:, None] * HC + + offs_c[None, :] + ).to(tl.float32) + + row_s = tl.load(ws_rs_ptr + (v_ws_base + t) * HC + offs_r).to(tl.float32) + if t == 0: + grad = grad / row_s[:, None] + row_corr = tl.sum(grad * (M - eps), axis=1) + else: + grad = grad / (row_s[:, None] + eps) + row_corr = tl.sum(grad * M, axis=1) + grad = grad - row_corr[:, None] + M = tl.load( + ws_M_ptr + M_ws_base + (2 * t) * HC * HC + offs_r[:, None] * HC + offs_c[None, :] + ).to(tl.float32) + + M_init = tl.load(M_init_ptr + mat_ptrs).to(tl.float32) + grad = grad * M_init + tl.store(grad_inp_ptr + mat_ptrs, grad.to(grad_inp_ptr.dtype.element_ty)) + + def _triton_sinkhorn_fwd( + input_logits: Tensor, num_iterations: int, eps: float = 1e-6 + ) -> Tuple[Tensor, Tensor]: + original_shape = input_logits.shape + hc = original_shape[-1] + N_batch = input_logits.numel() // (hc * hc) + dev = input_logits.device + out = torch.empty(N_batch, hc, hc, dtype=input_logits.dtype, device=dev) + M_init = torch.empty(N_batch, hc, hc, dtype=input_logits.dtype, device=dev) + inp = input_logits.contiguous().view(N_batch, hc, hc) + _triton_sinkhorn_fwd_kernel[(N_batch,)](inp, out, M_init, N_batch, eps, hc, num_iterations) + return out.view(original_shape), M_init.view(original_shape) + + def _triton_sinkhorn_bwd( + grad_output: Tensor, M_init: Tensor, num_iterations: int, eps: float = 1e-6 + ) -> Tensor: + original_shape = grad_output.shape + hc = original_shape[-1] + N_batch = grad_output.numel() // (hc * hc) + dev = grad_output.device + grad_input = torch.empty(N_batch, hc, hc, dtype=grad_output.dtype, device=dev) + go = grad_output.contiguous().view(N_batch, hc, hc) + mi = M_init.contiguous().view(N_batch, hc, hc) + ws_M = torch.empty(N_batch * 2 * num_iterations * hc * hc, dtype=torch.float32, device=dev) + ws_rs = torch.empty(N_batch * num_iterations * hc, dtype=torch.float32, device=dev) + ws_cs = torch.empty(N_batch * num_iterations * hc, dtype=torch.float32, device=dev) + _triton_sinkhorn_bwd_kernel[(N_batch,)]( + go, mi, grad_input, ws_M, ws_rs, ws_cs, N_batch, eps, hc, num_iterations + ) + return grad_input.view(original_shape) + + class TritonFusedSinkhorn(torch.autograd.Function): + """Autograd wrapper for Triton fused Sinkhorn.""" + + @staticmethod + def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): + """Run Triton Sinkhorn forward and save initial matrix for backward.""" + out, M_init = _triton_sinkhorn_fwd(input_logits, num_iterations, eps) + ctx.save_for_backward(M_init) + ctx.num_iterations = num_iterations + ctx.eps = eps + return out + + @staticmethod + def backward(ctx, grad_output: Tensor): + """Run Triton Sinkhorn backward.""" + (M_init,) = ctx.saved_tensors + grad_input = _triton_sinkhorn_bwd(grad_output, M_init, ctx.num_iterations, ctx.eps) + return grad_input, None, None + + def triton_fused_sinkhorn( + input_logits: Tensor, num_iterations: int, eps: float = 1e-6 + ) -> Tensor: + """Apply Triton fused Sinkhorn with autograd support.""" + return TritonFusedSinkhorn.apply(input_logits, num_iterations, eps) + + # ============================================================================ + # H_aggregate forward + # ============================================================================ + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) + for bc in (64, 128, 256, 512) + for bs in (1, 2, 4, 8) + for nw in (2, 4, 8) + ], + key=["C", "N"], + ) + @triton.jit + def _triton_h_agg_fwd_kernel( + x_ptr, + h_ptr, + out_ptr, + sb, + C: tl.constexpr, + N: tl.constexpr, + stride_x_s, + stride_x_n, + stride_x_c, + BLOCK_C: tl.constexpr, + BLOCK_S: tl.constexpr, + ): + """out[s, c] = sum_i x[s, i, c] * h[s, i].""" + pid_s = tl.program_id(0) + pid_c = tl.program_id(1) + offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + offs_c = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) + mask_s = offs_s < sb + mask_c = offs_c < C + mask_2d = mask_s[:, None] & mask_c[None, :] + + acc = tl.zeros((BLOCK_S, BLOCK_C), dtype=tl.float32) + for i in tl.static_range(N): + x_i = tl.load( + x_ptr + offs_s[:, None] * stride_x_s + i * stride_x_n + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + h_i = tl.load(h_ptr + offs_s * N + i, mask=mask_s, other=0.0).to(tl.float32) + acc += h_i[:, None] * x_i + tl.store( + out_ptr + offs_s[:, None] * C + offs_c[None, :], + acc.to(out_ptr.dtype.element_ty), + mask=mask_2d, + ) + + def _triton_h_aggregate_fwd(x: Tensor, h_pre: Tensor) -> Tensor: + s, b, n, C = x.shape + sb = s * b + out = torch.empty(sb, C, dtype=x.dtype, device=x.device) + x_flat = x.contiguous().view(sb, n, C) + h_flat = h_pre.contiguous().view(sb, n) + + grid = lambda META: (triton.cdiv(sb, META["BLOCK_S"]), triton.cdiv(C, META["BLOCK_C"])) + _triton_h_agg_fwd_kernel[grid]( + x_flat, h_flat, out, sb, C, n, x_flat.stride(0), x_flat.stride(1), x_flat.stride(2) + ) + return out.view(s, b, C) + + # ============================================================================ + # H_post BDA + # ============================================================================ + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) + for bc in (64, 128, 256, 512) + for bs in (1, 2, 4, 8) + for nw in (2, 4, 8) + ], + key=["C", "N"], + ) + @triton.jit + def _triton_hpb_fwd_kernel( + hr_ptr, + orig_ptr, + hp_ptr, + x_ptr, + bias_ptr, + out_ptr, + sb, + C: tl.constexpr, + N: tl.constexpr, + stride_hr_s, + stride_hr_i, + stride_hr_j, + stride_orig_s, + stride_orig_n, + stride_orig_c, + stride_out_s, + stride_out_n, + stride_out_c, + HAS_BIAS: tl.constexpr, + BLOCK_C: tl.constexpr, + BLOCK_S: tl.constexpr, + ): + """out = hr.T @ orig + hp * (x + bias).""" + pid_s = tl.program_id(0) + pid_c = tl.program_id(1) + offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + offs_c = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) + mask_s = offs_s < sb + mask_c = offs_c < C + mask_2d = mask_s[:, None] & mask_c[None, :] + + x_tile = tl.load(x_ptr + offs_s[:, None] * C + offs_c[None, :], mask=mask_2d, other=0.0).to( + tl.float32 + ) + if HAS_BIAS: + bias_tile = tl.load(bias_ptr + offs_c, mask=mask_c, other=0.0).to(tl.float32) + x_tile += bias_tile[None, :] + + for i in tl.static_range(N): + hp_i = tl.load(hp_ptr + offs_s * N + i, mask=mask_s, other=0.0).to(tl.float32) + out_i = hp_i[:, None] * x_tile + + for j in tl.static_range(N): + hr_ji = tl.load( + hr_ptr + offs_s * stride_hr_s + j * stride_hr_i + i * stride_hr_j, + mask=mask_s, + other=0.0, + ).to(tl.float32) + orig_j = tl.load( + orig_ptr + + offs_s[:, None] * stride_orig_s + + j * stride_orig_n + + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + out_i += hr_ji[:, None] * orig_j + + tl.store( + out_ptr + offs_s[:, None] * stride_out_s + i * stride_out_n + offs_c[None, :], + out_i.to(out_ptr.dtype.element_ty), + mask=mask_2d, + ) + + def _triton_h_post_bda_fwd( + h_res: Tensor, original_residual: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] + ) -> Tensor: + s, b, n, C = original_residual.shape + sb = s * b + dev = h_res.device + out = torch.empty(sb, n, C, dtype=h_res.dtype, device=dev) + hr_flat = h_res.contiguous().view(sb, n, n) + orig_flat = original_residual.contiguous().view(sb, n, C) + hp_flat = h_post.contiguous().view(sb, n) + x_flat = x.contiguous().view(sb, C) + + grid = lambda META: (triton.cdiv(sb, META["BLOCK_S"]), triton.cdiv(C, META["BLOCK_C"])) + _triton_hpb_fwd_kernel[grid]( + hr_flat, + orig_flat, + hp_flat, + x_flat, + bias if bias is not None else x_flat, + out, + sb, + C, + n, + hr_flat.stride(0), + hr_flat.stride(1), + hr_flat.stride(2), + orig_flat.stride(0), + orig_flat.stride(1), + orig_flat.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + HAS_BIAS=(bias is not None), + ) + return out.view(s, b, n, C) + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) + for bc in (64, 128, 256, 512) + for bs in (1, 2, 4, 8) + for nw in (2, 4, 8) + ], + key=["C", "N"], + ) + @triton.jit + def _triton_hpb_bwd_g_x_orig_kernel( + go_ptr, + hr_ptr, + hp_ptr, + g_orig_ptr, + g_x_ptr, + sb, + C: tl.constexpr, + N: tl.constexpr, + stride_go_s, + stride_go_n, + stride_go_c, + stride_hr_s, + stride_hr_i, + stride_hr_j, + stride_orig_s, + stride_orig_n, + stride_orig_c, + BLOCK_C: tl.constexpr, + BLOCK_S: tl.constexpr, + ): + """g_x = hp @ go, g_orig = hr @ go.""" + pid_s = tl.program_id(0) + pid_c = tl.program_id(1) + offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + offs_c = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) + mask_s = offs_s < sb + mask_c = offs_c < C + mask_2d = mask_s[:, None] & mask_c[None, :] + + g_x_acc = tl.zeros((BLOCK_S, BLOCK_C), dtype=tl.float32) + for j in tl.static_range(N): + go_j = tl.load( + go_ptr + offs_s[:, None] * stride_go_s + j * stride_go_n + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + hp_j = tl.load(hp_ptr + offs_s * N + j, mask=mask_s, other=0.0).to(tl.float32) + g_x_acc += hp_j[:, None] * go_j + tl.store( + g_x_ptr + offs_s[:, None] * C + offs_c[None, :], + g_x_acc.to(g_x_ptr.dtype.element_ty), + mask=mask_2d, + ) + + for i in tl.static_range(N): + g_orig_i = tl.zeros((BLOCK_S, BLOCK_C), dtype=tl.float32) + for j in tl.static_range(N): + go_j = tl.load( + go_ptr + offs_s[:, None] * stride_go_s + j * stride_go_n + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + hr_ij = tl.load( + hr_ptr + offs_s * stride_hr_s + i * stride_hr_i + j * stride_hr_j, + mask=mask_s, + other=0.0, + ).to(tl.float32) + g_orig_i += hr_ij[:, None] * go_j + tl.store( + g_orig_ptr + offs_s[:, None] * stride_orig_s + i * stride_orig_n + offs_c[None, :], + g_orig_i.to(g_orig_ptr.dtype.element_ty), + mask=mask_2d, + ) + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) + for bc in (64, 128, 256, 512) + for bs in (1, 2, 4, 8) + for nw in (2, 4, 8) + ], + key=["C", "N"], + ) + @triton.jit + def _triton_hpb_bwd_g_hp_hr_kernel( + go_ptr, + orig_ptr, + x_ptr, + bias_ptr, + g_hr_ptr, + g_hp_ptr, + sb, + C: tl.constexpr, + N: tl.constexpr, + stride_go_s, + stride_go_n, + stride_go_c, + stride_orig_s, + stride_orig_n, + stride_orig_c, + stride_hr_s, + stride_hr_i, + stride_hr_j, + HAS_BIAS: tl.constexpr, + BLOCK_C: tl.constexpr, + BLOCK_S: tl.constexpr, + ): + """g_hp = sum_c go*(x+bias), g_hr = orig @ go.T.""" + pid_s = tl.program_id(0) + offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + mask_s = offs_s < sb + + g_hp_acc = tl.zeros((BLOCK_S, N), dtype=tl.float32) + g_hr_acc = tl.zeros((BLOCK_S, N * N), dtype=tl.float32) + + for c_start in range(0, C, BLOCK_C): + offs_c = c_start + tl.arange(0, BLOCK_C) + mask_c = offs_c < C + mask_2d = mask_s[:, None] & mask_c[None, :] + + x_tile = tl.load( + x_ptr + offs_s[:, None] * C + offs_c[None, :], mask=mask_2d, other=0.0 + ).to(tl.float32) + if HAS_BIAS: + bias_tile = tl.load(bias_ptr + offs_c, mask=mask_c, other=0.0).to(tl.float32) + x_tile += bias_tile[None, :] + + for i in tl.static_range(N): + go_i = tl.load( + go_ptr + offs_s[:, None] * stride_go_s + i * stride_go_n + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + dot_hp = tl.sum(go_i * x_tile, axis=1) + g_hp_acc += tl.where( + tl.arange(0, N)[None, :] == i, + dot_hp[:, None], + tl.zeros((BLOCK_S, N), dtype=tl.float32), + ) + for j in tl.static_range(N): + orig_j = tl.load( + orig_ptr + + offs_s[:, None] * stride_orig_s + + j * stride_orig_n + + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + dot_hr = tl.sum(go_i * orig_j, axis=1) + g_hr_acc += tl.where( + tl.arange(0, N * N)[None, :] == j * N + i, + dot_hr[:, None], + tl.zeros((BLOCK_S, N * N), dtype=tl.float32), + ) + + offs_n = tl.arange(0, N) + tl.store( + g_hp_ptr + offs_s[:, None] * N + offs_n[None, :], + g_hp_acc.to(g_hp_ptr.dtype.element_ty), + mask=mask_s[:, None], + ) + + # N is expected to stay small for mHC, so this simple extraction is acceptable. + nn_offs = tl.arange(0, N * N) + for i in tl.static_range(N): + for j in tl.static_range(N): + col_mask = (nn_offs == (i * N + j)).to(tl.float32) + val = tl.sum(g_hr_acc * col_mask[None, :], axis=1) + tl.store( + g_hr_ptr + offs_s * stride_hr_s + i * stride_hr_i + j * stride_hr_j, + val.to(g_hr_ptr.dtype.element_ty), + mask=mask_s, + ) + + def _triton_h_post_bda_bwd( + grad_output: Tensor, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + x: Tensor, + bias: Optional[Tensor], + ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Optional[Tensor]]: + s, b, n, C = original_residual.shape + sb = s * b + dev = h_res.device + + g_hr = torch.empty(sb, n, n, dtype=h_res.dtype, device=dev) + g_res = torch.empty(sb, n, C, dtype=original_residual.dtype, device=dev) + g_hp = torch.empty(sb, n, dtype=h_post.dtype, device=dev) + g_x = torch.empty(sb, C, dtype=x.dtype, device=dev) + + go_flat = grad_output.contiguous().view(sb, n, C) + hr_flat = h_res.contiguous().view(sb, n, n) + orig_flat = original_residual.contiguous().view(sb, n, C) + hp_flat = h_post.contiguous().view(sb, n) + x_flat = x.contiguous().view(sb, C) + + grid_a = lambda META: (triton.cdiv(sb, META["BLOCK_S"]), triton.cdiv(C, META["BLOCK_C"])) + _triton_hpb_bwd_g_x_orig_kernel[grid_a]( + go_flat, + hr_flat, + hp_flat, + g_res, + g_x, + sb, + C, + n, + go_flat.stride(0), + go_flat.stride(1), + go_flat.stride(2), + hr_flat.stride(0), + hr_flat.stride(1), + hr_flat.stride(2), + g_res.stride(0), + g_res.stride(1), + g_res.stride(2), + ) + + grid_b = lambda META: (triton.cdiv(sb, META["BLOCK_S"]),) + _triton_hpb_bwd_g_hp_hr_kernel[grid_b]( + go_flat, + orig_flat, + x_flat, + bias if bias is not None else x_flat, + g_hr, + g_hp, + sb, + C, + n, + go_flat.stride(0), + go_flat.stride(1), + go_flat.stride(2), + orig_flat.stride(0), + orig_flat.stride(1), + orig_flat.stride(2), + g_hr.stride(0), + g_hr.stride(1), + g_hr.stride(2), + HAS_BIAS=(bias is not None), + ) + + g_bias = g_x.sum(dim=0).to(dtype=bias.dtype) if bias is not None else None + return ( + g_hr.view(s, b, n, n), + g_res.view(s, b, n, C), + g_hp.view(s, b, n), + g_x.view(s, b, C), + g_bias, + ) + + +_TRITON_IMPLS = ( + { + "sinkhorn": triton_fused_sinkhorn, + "h_aggregate_fwd": _triton_h_aggregate_fwd, + "h_post_bda_fwd": _triton_h_post_bda_fwd, + "h_post_bda_bwd": _triton_h_post_bda_bwd, + } + if _TRITON_AVAILABLE + else {"sinkhorn": None, "h_aggregate_fwd": None, "h_post_bda_fwd": None, "h_post_bda_bwd": None} +) + + +# ============================================================================ +# CuTile implementations (only defined when cuda.tile is available) +# ============================================================================ + +if _CUTILE_AVAILABLE: + ConstInt = ct.Constant[int] + PAD_ZERO = ct.PaddingMode.ZERO + + # -- Sinkhorn kernels ---------------------------------------------------- + + @ct.kernel + def _ct_sinkhorn_fwd_kernel( + inp, out, M_init_out, eps, HC: ConstInt, NUM_ITERS: ConstInt, TILE_SIZE: ConstInt + ): + pid = ct.bid(0) + logits = ct.load(inp, index=(pid, 0, 0), shape=(TILE_SIZE, HC, HC)).astype(ct.float32) + row_max = ct.max(logits, axis=2, keepdims=True) + M = ct.exp2((logits - row_max) * LOG2E) + ct.store( + M_init_out, + index=(pid, 0, 0), + tile=ct.reshape(M.astype(M_init_out.dtype), (TILE_SIZE, HC, HC)), + ) + row_sum = ct.sum(M, axis=2, keepdims=True) + M = M / row_sum + eps + col_sum = ct.sum(M, axis=1, keepdims=True) + M = M / (col_sum + eps) + for _ in range(NUM_ITERS - 1): + row_sum = ct.sum(M, axis=2, keepdims=True) + M = M / (row_sum + eps) + col_sum = ct.sum(M, axis=1, keepdims=True) + M = M / (col_sum + eps) + ct.store(out, index=(pid, 0, 0), tile=ct.reshape(M.astype(out.dtype), (TILE_SIZE, HC, HC))) + + @ct.kernel + def _ct_sinkhorn_bwd_kernel( + grad_out, + M_init, + grad_inp, + ws_M, + ws_rs, + ws_cs, + eps, + HC: ConstInt, + NUM_ITERS: ConstInt, + TILE_SIZE: ConstInt, + ): + pid = ct.bid(0) + M_base = pid * (2 * NUM_ITERS) + v_base = pid * NUM_ITERS + + M = ct.load(M_init, index=(pid, 0, 0), shape=(TILE_SIZE, HC, HC)).astype(ct.float32) + for t in range(NUM_ITERS): + ct.store(ws_M, index=(M_base + 2 * t, 0, 0), tile=M) + row_sum = ct.sum(M, axis=2, keepdims=True) + ct.store(ws_rs, index=(v_base + t, 0, 0), tile=row_sum) + if t == 0: + M = M / row_sum + eps + else: + M = M / (row_sum + eps) + ct.store(ws_M, index=(M_base + 2 * t + 1, 0, 0), tile=M) + col_sum = ct.sum(M, axis=1, keepdims=True) + ct.store(ws_cs, index=(v_base + t, 0, 0), tile=col_sum) + M = M / (col_sum + eps) + + grad = ct.load(grad_out, index=(pid, 0, 0), shape=(TILE_SIZE, HC, HC)).astype(ct.float32) + for t_rev in range(NUM_ITERS): + t = NUM_ITERS - 1 - t_rev + col_s = ct.load(ws_cs, index=(v_base + t, 0, 0), shape=(TILE_SIZE, 1, HC)) + grad = grad / (col_s + eps) + col_corr = ct.sum(grad * M, axis=1, keepdims=True) + grad = grad - col_corr + M = ct.load(ws_M, index=(M_base + 2 * t + 1, 0, 0), shape=(TILE_SIZE, HC, HC)) + row_s = ct.load(ws_rs, index=(v_base + t, 0, 0), shape=(TILE_SIZE, HC, 1)) + if t == 0: + grad = grad / row_s + row_corr = ct.sum(grad * (M - eps), axis=2, keepdims=True) + else: + grad = grad / (row_s + eps) + row_corr = ct.sum(grad * M, axis=2, keepdims=True) + grad = grad - row_corr + M = ct.load(ws_M, index=(M_base + 2 * t, 0, 0), shape=(TILE_SIZE, HC, HC)) + grad = grad * M + ct.store(grad_inp, index=(pid, 0, 0), tile=grad.astype(grad_inp.dtype)) + + def _sinkhorn_autotune_tile_sizes(N_batch): + """Generate autotune search space for sinkhorn kernels.""" + for ts in (1, 2, 4, 8, 16, 32, 64, 128): + if ts <= N_batch: + yield ts + + _sinkhorn_fwd_best_cfg: dict = {} + _sinkhorn_bwd_best_cfg: dict = {} + + def _cutile_sinkhorn_fwd( + input_logits: Tensor, num_iterations: int, eps: float = 1e-6 + ) -> Tuple[Tensor, Tensor]: + original_shape = input_logits.shape + hc = original_shape[-1] + N_batch = input_logits.numel() // (hc * hc) + dev = input_logits.device + stream = torch.cuda.current_stream() + out = torch.empty(N_batch, hc, hc, dtype=input_logits.dtype, device=dev) + M_init = torch.empty(N_batch, hc, hc, dtype=input_logits.dtype, device=dev) + inp = input_logits.view(N_batch, hc, hc) + + cache_key = (N_batch, hc, num_iterations) + cached = _sinkhorn_fwd_best_cfg.get(cache_key) + + if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + ts = cached if cached is not None else math.gcd(N_batch, 128) + ct.launch( + stream, + (math.ceil(N_batch / ts), 1, 1), + _ct_sinkhorn_fwd_kernel, + (inp, out, M_init, eps, hc, num_iterations, ts), + ) + else: + from types import SimpleNamespace + + configs = [ + SimpleNamespace(TILE_SIZE=ts) for ts in _sinkhorn_autotune_tile_sizes(N_batch) + ] + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(N_batch / cfg.TILE_SIZE), 1, 1), + kernel=_ct_sinkhorn_fwd_kernel, + args_fn=lambda cfg: (inp, out, M_init, eps, hc, num_iterations, cfg.TILE_SIZE), + search_space=configs, + ) + best_ts = tuned.tuned_config.TILE_SIZE + _sinkhorn_fwd_best_cfg[cache_key] = best_ts + ct.launch( + stream, + (math.ceil(N_batch / best_ts), 1, 1), + _ct_sinkhorn_fwd_kernel, + (inp, out, M_init, eps, hc, num_iterations, best_ts), + ) + + return out.view(original_shape), M_init.view(original_shape) + + def _cutile_sinkhorn_bwd( + grad_output: Tensor, M_init: Tensor, num_iterations: int, eps: float = 1e-6 + ) -> Tensor: + original_shape = grad_output.shape + hc = original_shape[-1] + N_batch = grad_output.numel() // (hc * hc) + dev = grad_output.device + stream = torch.cuda.current_stream() + grad_input = torch.empty(N_batch, hc, hc, dtype=grad_output.dtype, device=dev) + go = grad_output.view(N_batch, hc, hc) + mi = M_init.view(N_batch, hc, hc) + + cache_key = (N_batch, hc, num_iterations) + cached = _sinkhorn_bwd_best_cfg.get(cache_key) + + def _alloc_and_launch(ts): + ws_M = torch.empty( + N_batch * 2 * num_iterations, hc, hc, dtype=torch.float32, device=dev + ) + ws_rs = torch.empty(N_batch * num_iterations, hc, 1, dtype=torch.float32, device=dev) + ws_cs = torch.empty(N_batch * num_iterations, 1, hc, dtype=torch.float32, device=dev) + ct.launch( + stream, + (math.ceil(N_batch / ts), 1, 1), + _ct_sinkhorn_bwd_kernel, + (go, mi, grad_input, ws_M, ws_rs, ws_cs, eps, hc, num_iterations, ts), + ) + + if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + ts = cached if cached is not None else math.gcd(N_batch, 128) + _alloc_and_launch(ts) + else: + from types import SimpleNamespace + + configs = [ + SimpleNamespace(TILE_SIZE=ts) for ts in _sinkhorn_autotune_tile_sizes(N_batch) + ] + # Allocate workspace for largest tile size (all configs share same workspace shape). + ws_M = torch.empty( + N_batch * 2 * num_iterations, hc, hc, dtype=torch.float32, device=dev + ) + ws_rs = torch.empty(N_batch * num_iterations, hc, 1, dtype=torch.float32, device=dev) + ws_cs = torch.empty(N_batch * num_iterations, 1, hc, dtype=torch.float32, device=dev) + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(N_batch / cfg.TILE_SIZE), 1, 1), + kernel=_ct_sinkhorn_bwd_kernel, + args_fn=lambda cfg: ( + go, + mi, + grad_input, + ws_M, + ws_rs, + ws_cs, + eps, + hc, + num_iterations, + cfg.TILE_SIZE, + ), + search_space=configs, + ) + best_ts = tuned.tuned_config.TILE_SIZE + _sinkhorn_bwd_best_cfg[cache_key] = best_ts + # Re-launch with best config. + _alloc_and_launch(best_ts) + + return grad_input.view(original_shape) + + # -- H_aggregate kernels ------------------------------------------------- + + @ct.kernel + def _ct_h_agg_fwd_kernel(x, h_pre, out, N: ConstInt, TILE_M: ConstInt, TILE_C: ConstInt): + pid = ct.bid(0) + num_tiles = ct.num_tiles(x, axis=2, shape=(TILE_M, N, TILE_C)) + h_tile = ct.load(h_pre, index=(pid, 0), shape=(TILE_M, N), padding_mode=PAD_ZERO) + h_tile = ct.expand_dims(h_tile, axis=2) + for j in range(num_tiles): + x_tile = ct.load(x, index=(pid, 0, j), shape=(TILE_M, N, TILE_C), padding_mode=PAD_ZERO) + acc = ct.sum(x_tile * h_tile, axis=1).astype(ct.float32) + ct.store(out, index=(pid, j), tile=acc.astype(out.dtype)) + + @ct.kernel + def _ct_h_agg_bwd_kernel(go, x, h_pre, gx, gh, N: ConstInt, TILE_M: ConstInt, TILE_C: ConstInt): + pid = ct.bid(0) + num_c_tiles = ct.num_tiles(go, axis=1, shape=(TILE_M, TILE_C)) + h_tile = ct.load(h_pre, index=(pid, 0), shape=(TILE_M, N), padding_mode=PAD_ZERO) + h_expanded = ct.expand_dims(h_tile, axis=2) + gh_acc = ct.full((TILE_M, N), 0, dtype=ct.float32) + for ct_idx in range(num_c_tiles): + go_tile = ct.load( + go, index=(pid, ct_idx), shape=(TILE_M, TILE_C), padding_mode=PAD_ZERO + ) + go_expanded = ct.expand_dims(go_tile, axis=1) + x_tile = ct.load( + x, index=(pid, 0, ct_idx), shape=(TILE_M, N, TILE_C), padding_mode=PAD_ZERO + ) + gx_tile = go_expanded * h_expanded + ct.store(gx, index=(pid, 0, ct_idx), tile=gx_tile.astype(gx.dtype)) + gh_acc += ct.sum(go_expanded * x_tile, axis=2) + ct.store(gh, index=(pid, 0), tile=gh_acc.astype(gh.dtype)) + + def _cutile_h_aggregate_fwd(x: Tensor, h_pre: Tensor) -> Tensor: + s, b, n, C = x.shape + sb = s * b + stream = torch.cuda.current_stream() + out = torch.empty(sb, C, dtype=x.dtype, device=x.device) + x_flat = x.view(sb, n, C) + h_flat = h_pre.view(sb, n) + + # Autotune disabled — causes cudaErrorLaunchFailure during training. + tm, tc = math.gcd(sb, 4), math.gcd(C, 1024) + ct.launch( + stream, (math.ceil(sb / tm),), _ct_h_agg_fwd_kernel, (x_flat, h_flat, out, n, tm, tc) + ) + + return out.view(s, b, C) + + def _cutile_h_aggregate_bwd( + grad_output: Tensor, x: Tensor, h_pre: Tensor + ) -> Tuple[Tensor, Tensor]: + s, b, n, C = x.shape + sb = s * b + stream = torch.cuda.current_stream() + gx = torch.empty(sb, n, C, dtype=x.dtype, device=x.device) + gh = torch.empty(sb, n, dtype=h_pre.dtype, device=x.device) + go_flat = grad_output.view(sb, C) + x_flat = x.view(sb, n, C) + h_flat = h_pre.view(sb, n) + + # Autotune disabled — causes cudaErrorLaunchFailure during training. + tm, tc = math.gcd(sb, 4), math.gcd(C, 1024) + ct.launch( + stream, + (math.ceil(sb / tm),), + _ct_h_agg_bwd_kernel, + (go_flat, x_flat, h_flat, gx, gh, n, tm, tc), + ) + + return gx.view(s, b, n, C), gh.view(s, b, n) + + # -- H_post BDA kernels -------------------------------------------------- + + @ct.kernel + def _ct_hpb_fwd_kernel( + hr, orig, hp, x, out, N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt + ): + pid = ct.bid(0) + num_c_tiles = ct.num_tiles(x, axis=1, shape=(TILE_SIZE, TILE_C)) + hp_tile = ct.load(hp, index=(pid, 0), shape=(TILE_SIZE, N), padding_mode=PAD_ZERO) + hp_exp = ct.expand_dims(hp_tile, axis=2) # (TILE_SIZE, N, 1) + hr_tile = ct.load(hr, index=(pid, 0, 0), shape=(TILE_SIZE, N, N), padding_mode=PAD_ZERO) + for ct_idx in range(num_c_tiles): + orig_tile = ct.load( + orig, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + x_tile = ct.load( + x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO + ) + x_exp = ct.expand_dims(x_tile, axis=1) # (TILE_SIZE, 1, TILE_C) + out_tile = hp_exp * x_exp # (TILE_SIZE, N, TILE_C) + for j in range(N): + hr_row = ct.extract(hr_tile, (0, j, 0), shape=(TILE_SIZE, 1, N)) + hr_col = ct.reshape(hr_row, (TILE_SIZE, N, 1)) + orig_row = ct.extract(orig_tile, (0, j, 0), shape=(TILE_SIZE, 1, TILE_C)) + out_tile = out_tile + hr_col * orig_row + ct.store(out, index=(pid, 0, ct_idx), tile=out_tile.astype(out.dtype)) + + @ct.kernel + def _ct_hpb_fwd_bias_kernel( + hr, orig, hp, x, bias, out, N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt + ): + pid = ct.bid(0) + num_c_tiles = ct.num_tiles(x, axis=1, shape=(TILE_SIZE, TILE_C)) + hp_tile = ct.load(hp, index=(pid, 0), shape=(TILE_SIZE, N), padding_mode=PAD_ZERO) + hp_exp = ct.expand_dims(hp_tile, axis=2) # (TILE_SIZE, N, 1) + hr_tile = ct.load(hr, index=(pid, 0, 0), shape=(TILE_SIZE, N, N), padding_mode=PAD_ZERO) + for ct_idx in range(num_c_tiles): + orig_tile = ct.load( + orig, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + x_tile = ct.load( + x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO + ) + bias_tile = ct.load(bias, index=(ct_idx,), shape=(TILE_C,), padding_mode=PAD_ZERO) + xb_exp = ct.expand_dims(x_tile + bias_tile, axis=1) # (TILE_SIZE, 1, TILE_C) + out_tile = hp_exp * xb_exp # (TILE_SIZE, N, TILE_C) + for j in range(N): + hr_row = ct.extract(hr_tile, (0, j, 0), shape=(TILE_SIZE, 1, N)) + hr_col = ct.reshape(hr_row, (TILE_SIZE, N, 1)) + orig_row = ct.extract(orig_tile, (0, j, 0), shape=(TILE_SIZE, 1, TILE_C)) + out_tile = out_tile + hr_col * orig_row + ct.store(out, index=(pid, 0, ct_idx), tile=out_tile.astype(out.dtype)) + + @ct.kernel + def _ct_hpb_bwd_g_x_orig_kernel( + go, hr, hp, g_orig, g_x, N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt + ): + """Compute g_x = hp @ go and g_orig = hr @ go. + + Grid: (ceil(sb / TILE_SIZE), ceil(C / TILE_C)). + 2D grid — no loop, no accumulators. + """ + pid = ct.bid(0) + ct_idx = ct.bid(1) + hp_tile = ct.load(hp, index=(pid, 0), shape=(TILE_SIZE, N), padding_mode=PAD_ZERO) + hr_tile = ct.load(hr, index=(pid, 0, 0), shape=(TILE_SIZE, N, N), padding_mode=PAD_ZERO) + go_tile = ct.load( + go, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + g_x_tile = ct.full((TILE_SIZE, 1, TILE_C), 0, dtype=ct.float32) + g_orig_tile = ct.full((TILE_SIZE, N, TILE_C), 0, dtype=ct.float32) + for j in range(N): + hp_j = ct.extract(hp_tile, (0, j), shape=(TILE_SIZE, 1)) + hp_j_exp = ct.expand_dims(hp_j, axis=2) # [TS, 1, 1] + go_j = ct.extract(go_tile, (0, j, 0), shape=(TILE_SIZE, 1, TILE_C)) + g_x_tile = g_x_tile + hp_j_exp * go_j + hr_col_j = ct.extract(hr_tile, (0, 0, j), shape=(TILE_SIZE, N, 1)) + g_orig_tile = g_orig_tile + hr_col_j * go_j + ct.store( + g_x, + index=(pid, ct_idx), + tile=ct.reshape(g_x_tile, (TILE_SIZE, TILE_C)).astype(g_x.dtype), + ) + ct.store(g_orig, index=(pid, 0, ct_idx), tile=g_orig_tile.astype(g_orig.dtype)) + + @ct.kernel + def _ct_hpb_bwd_g_hp_hr_kernel( + go, orig, x, g_hr, g_hp, N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt + ): + """Compute g_hp = sum(go * x) and g_hr = orig @ go.T (no bias). + + Grid: (ceil(sb / TILE_SIZE),). Loops over C-tiles. + """ + pid = ct.bid(0) + num_c_tiles = ct.cdiv(go.shape[2], TILE_C) + acc_g_hp = ct.full((TILE_SIZE, N, 1), 0, dtype=ct.float32) + acc_g_hr = ct.full((TILE_SIZE, N, N), 0, dtype=ct.float32) + for ct_idx in range(num_c_tiles): + x_tile = ct.load( + x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO + ) + x_exp = ct.expand_dims(x_tile, axis=1) # [TS, 1, TC] + go_tile = ct.load( + go, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + orig_tile = ct.load( + orig, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + acc_g_hp = acc_g_hp + ct.sum(go_tile * x_exp, axis=2, keepdims=True) + acc_g_hr = acc_g_hr + ct.sum( + ct.expand_dims(orig_tile, axis=2) * ct.expand_dims(go_tile, axis=1), axis=3 + ) + ct.store(g_hp, index=(pid, 0), tile=ct.reshape(acc_g_hp, (TILE_SIZE, N)).astype(g_hp.dtype)) + ct.store(g_hr, index=(pid, 0, 0), tile=acc_g_hr.astype(g_hr.dtype)) + + @ct.kernel + def _ct_hpb_bwd_g_hp_hr_bias_kernel( + go, orig, x, bias, g_hr, g_hp, N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt + ): + """Compute g_hp = sum(go * (x+bias)) and g_hr = orig @ go.T (with bias). + + Grid: (ceil(sb / TILE_SIZE),). Loops over C-tiles. + """ + pid = ct.bid(0) + num_c_tiles = ct.cdiv(go.shape[2], TILE_C) + acc_g_hp = ct.full((TILE_SIZE, N, 1), 0, dtype=ct.float32) + acc_g_hr = ct.full((TILE_SIZE, N, N), 0, dtype=ct.float32) + for ct_idx in range(num_c_tiles): + x_tile = ct.load( + x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO + ) + bias_tile = ct.load(bias, index=(ct_idx,), shape=(TILE_C,), padding_mode=PAD_ZERO) + xb_exp = ct.expand_dims(x_tile + bias_tile, axis=1) # [TS, 1, TC] + go_tile = ct.load( + go, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + orig_tile = ct.load( + orig, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + acc_g_hp = acc_g_hp + ct.sum(go_tile * xb_exp, axis=2, keepdims=True) + acc_g_hr = acc_g_hr + ct.sum( + ct.expand_dims(orig_tile, axis=2) * ct.expand_dims(go_tile, axis=1), axis=3 + ) + ct.store(g_hp, index=(pid, 0), tile=ct.reshape(acc_g_hp, (TILE_SIZE, N)).astype(g_hp.dtype)) + ct.store(g_hr, index=(pid, 0, 0), tile=acc_g_hr.astype(g_hr.dtype)) + + # -- H_post BDA autotune configs & caches -------------------------------- + + def _hpb_autotune_configs(sb, C): + """Generate TILE_SIZE × TILE_C search space for h_post_bda kernels.""" + for tile_size in (1, 2, 4, 8): + for tile_c in (32, 64, 128, 256, 512, 1024): + if tile_c <= C and tile_size <= sb: + yield {"TILE_SIZE": tile_size, "TILE_C": tile_c} + + _hpb_fwd_best_cfg: dict = {} + _hpb_bwd_g_x_orig_best_cfg: dict = {} + _hpb_bwd_g_hp_hr_best_cfg: dict = {} + + def _cutile_h_post_bda_fwd( + h_res: Tensor, original_residual: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] + ) -> Tensor: + s, b, n, C = original_residual.shape + sb = s * b + stream = torch.cuda.current_stream() + out = torch.empty(sb, n, C, dtype=h_res.dtype, device=h_res.device) + hr_flat = h_res.view(sb, n, n) + orig_flat = original_residual.view(sb, n, C) + hp_flat = h_post.view(sb, n) + x_flat = x.view(sb, C) + + cache_key = (sb, n, C, bias is not None) + cached = _hpb_fwd_best_cfg.get(cache_key) + kernel = _ct_hpb_fwd_bias_kernel if bias is not None else _ct_hpb_fwd_kernel + + # Autotune disabled — causes cudaErrorLaunchFailure during training. + if cached is not None: + ts, tc = cached + else: + ts, tc = 1, math.gcd(C, 1024) + args = (hr_flat, orig_flat, hp_flat, x_flat) + if bias is not None: + args = args + (bias,) + args = args + (out, n, tc, ts) + ct.launch(stream, (math.ceil(sb / ts),), kernel, args) + + return out.view(s, b, n, C) + + def _cutile_h_post_bda_bwd( + grad_output: Tensor, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + x: Tensor, + bias: Optional[Tensor], + ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Optional[Tensor]]: + s, b, n, C = original_residual.shape + sb = s * b + stream = torch.cuda.current_stream() + g_hr = torch.empty(sb, n, n, dtype=h_res.dtype, device=h_res.device) + g_res = torch.empty(sb, n, C, dtype=original_residual.dtype, device=h_res.device) + g_hp = torch.empty(sb, n, dtype=h_post.dtype, device=h_res.device) + g_x = torch.empty(sb, C, dtype=x.dtype, device=h_res.device) + go_flat = grad_output.view(sb, n, C) + hr_flat = h_res.view(sb, n, n) + orig_flat = original_residual.view(sb, n, C) + hp_flat = h_post.view(sb, n) + x_flat = x.view(sb, C) + + # --- Kernel A: g_x, g_orig (2D grid, no loop) --- + cache_key_a = ('hpb_bwd_g_x_orig', sb, n, C) + cached_a = _hpb_bwd_g_x_orig_best_cfg.get(cache_key_a) + + if cached_a is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + if cached_a is not None: + ts, tc = cached_a + else: + ts, tc = 1, math.gcd(C, 1024) + ct.launch( + stream, + (math.ceil(sb / ts), math.ceil(C / tc)), + _ct_hpb_bwd_g_x_orig_kernel, + (go_flat, hr_flat, hp_flat, g_res, g_x, n, tc, ts), + ) + else: + from types import SimpleNamespace + + configs = [SimpleNamespace(**c) for c in _hpb_autotune_configs(sb, C)] + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(sb / cfg.TILE_SIZE), math.ceil(C / cfg.TILE_C)), + kernel=_ct_hpb_bwd_g_x_orig_kernel, + args_fn=lambda cfg: ( + go_flat, + hr_flat, + hp_flat, + g_res, + g_x, + n, + cfg.TILE_C, + cfg.TILE_SIZE, + ), + search_space=configs, + ) + best = tuned.tuned_config + _hpb_bwd_g_x_orig_best_cfg[cache_key_a] = (best.TILE_SIZE, best.TILE_C) + ct.launch( + stream, + (math.ceil(sb / best.TILE_SIZE), math.ceil(C / best.TILE_C)), + _ct_hpb_bwd_g_x_orig_kernel, + (go_flat, hr_flat, hp_flat, g_res, g_x, n, best.TILE_C, best.TILE_SIZE), + ) + + # --- Kernel B: g_hp, g_hr (1D grid, loops C-tiles) --- + cache_key_b = ('hpb_bwd_g_hp_hr', sb, n, C, bias is not None) + cached_b = _hpb_bwd_g_hp_hr_best_cfg.get(cache_key_b) + hp_hr_kernel = ( + _ct_hpb_bwd_g_hp_hr_bias_kernel if bias is not None else _ct_hpb_bwd_g_hp_hr_kernel + ) + + if cached_b is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + if cached_b is not None: + ts, tc = cached_b + else: + ts, tc = 1, math.gcd(C, 1024) + args = (go_flat, orig_flat, x_flat) + if bias is not None: + args = args + (bias,) + args = args + (g_hr, g_hp, n, tc, ts) + ct.launch(stream, (math.ceil(sb / ts),), hp_hr_kernel, args) + else: + from types import SimpleNamespace + + configs = [SimpleNamespace(**c) for c in _hpb_autotune_configs(sb, C)] + + def _hp_hr_args_fn(cfg): + args = (go_flat, orig_flat, x_flat) + if bias is not None: + args = args + (bias,) + return args + (g_hr, g_hp, n, cfg.TILE_C, cfg.TILE_SIZE) + + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(sb / cfg.TILE_SIZE),), + kernel=hp_hr_kernel, + args_fn=_hp_hr_args_fn, + search_space=configs, + ) + best = tuned.tuned_config + _hpb_bwd_g_hp_hr_best_cfg[cache_key_b] = (best.TILE_SIZE, best.TILE_C) + args = (go_flat, orig_flat, x_flat) + if bias is not None: + args = args + (bias,) + args = args + (g_hr, g_hp, n, best.TILE_C, best.TILE_SIZE) + ct.launch(stream, (math.ceil(sb / best.TILE_SIZE),), hp_hr_kernel, args) + + g_bias = g_x.sum(dim=0).to(dtype=bias.dtype) if bias is not None else None + return ( + g_hr.view(s, b, n, n), + g_res.view(s, b, n, C), + g_hp.view(s, b, n), + g_x.view(s, b, C), + g_bias, + ) + + # -- Proj RMS kernels ---------------------------------------------------- + + @ct.function + def _ct_rms_dnorm(a_tile, norm_tile, dr_tile, K, eps=1e-6): + inv_norm = ct.where(norm_tile > 0, 1.0 / norm_tile, 0.0) + inv_sqrt_k = 1.0 / ct.sqrt(K) + u = norm_tile * inv_sqrt_k + eps + coeff = -(1.0 / (u * u)) * inv_sqrt_k + return dr_tile * coeff * a_tile * inv_norm + + @ct.kernel + def _ct_proj_rms_fwd_kernel( + A, + B, + PROJ, + NORM, + R, + M: int, + N: int, + K: int, + eps: float, + TILE_M: ConstInt, + TILE_N: ConstInt, + TILE_K: ConstInt, + SPLIT_K: ConstInt, + ): + ''' + Grid: (num_tiles_m, num_tiles_k). + Fused matmul + norm + r: proj, norm, r in one pass over K. + R is a retained signature placeholder; r is computed after split-K + reduction from NORM. + ''' + tile_m_id = ct.bid(0) + split_k_id = ct.bid(1) + num_m_tiles = ct.cdiv(M, TILE_M) + num_k_tiles = ct.cdiv(K, TILE_K) + num_k_tiles_per_split = ct.cdiv(num_k_tiles, SPLIT_K) + tile_k_id_start = split_k_id * num_k_tiles_per_split + tile_k_id_end = ct.minimum(tile_k_id_start + num_k_tiles_per_split, num_k_tiles) + acc = ct.full((TILE_M, TILE_N), 0.0, dtype=ct.float32) + sum_sq = ct.full((TILE_M, 1), 0.0, dtype=ct.float32) + for tile_k_id in range(tile_k_id_start, tile_k_id_end): + a_tile = ct.load( + A, index=(tile_m_id, tile_k_id), shape=(TILE_M, TILE_K), padding_mode=PAD_ZERO + ) + b_tile = ct.load(B, index=(0, tile_k_id), shape=(TILE_N, TILE_K), padding_mode=PAD_ZERO) + acc = ct.mma( + a_tile.astype(ct.tfloat32), b_tile.transpose().astype(ct.tfloat32), acc=acc + ) + sum_sq += ct.sum(a_tile * a_tile, axis=1, keepdims=True) + + bid_m_k = tile_m_id + split_k_id * num_m_tiles + ct.store(PROJ, index=(bid_m_k, 0), tile=acc.astype(PROJ.dtype)) + ct.store(NORM, index=(bid_m_k, 0), tile=sum_sq.astype(NORM.dtype)) + + # -- Sigmoid helper for cuTile kernels ------------------------------------ + + @ct.function + def _ct_sigmoid(x): + """Sigmoid via exp2: σ(x) = 1 / (1 + 2^(-x * log2(e))).""" + return 1.0 / (1.0 + ct.exp2(-x * LOG2E)) + + # -- Reduce split-K + compute_h kernel ------------------------------------ + + @ct.kernel + def _ct_reduce_compute_h_kernel( + Y_acc, + R_acc, + Bias, + Alpha_pre, + Alpha_post, + Alpha_res, + H_PRE, + H_POST, + H_RES, + R, + PROJ_OUT, + M: int, + N: int, + K: int, + n: ConstInt, + eps: float, + compute_h_eps: float, + TILE_SIZE_M: ConstInt, + TILE_SIZE_N: ConstInt, + SPLIT_K: ConstInt, + ): + """Reduce split-K partial proj/norm, compute r, and apply compute_h activations. + + Grid: (ceil(M / TILE_SIZE_M),). + TILE_SIZE_N = next_power_of_2(N) so one tile covers the full N dimension. + Alpha_{pre,post,res} are [1] tensors (scalar parameters). + """ + bid_m = ct.bid(0) + num_bid_m = ct.cdiv(M, TILE_SIZE_M) + + alpha_pre = ct.load(Alpha_pre, index=(0,), shape=(1,)).item() + alpha_post = ct.load(Alpha_post, index=(0,), shape=(1,)).item() + alpha_res = ct.load(Alpha_res, index=(0,), shape=(1,)).item() + + # 1. Reduce split-K partials for each logical output segment. + pre_accum = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) + post_accum = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) + r_accum = ct.full((TILE_SIZE_M, 1), 0.0, dtype=ct.float32) + + for split_idx in ct.static_iter(range(SPLIT_K)): + bid_m_k = bid_m + split_idx * num_bid_m + pre_tile = ct.load( + Y_acc, index=(bid_m_k, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO + ) + post_tile = ct.load( + Y_acc, index=(bid_m_k, 1), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO + ) + pre_accum = pre_accum + ct.astype(pre_tile, ct.float32) + post_accum = post_accum + ct.astype(post_tile, ct.float32) + + r_tile = ct.load( + R_acc, index=(bid_m_k, 0), shape=(TILE_SIZE_M, 1), padding_mode=PAD_ZERO + ) + r_accum = r_accum + ct.astype(r_tile, ct.float32) + + # Store reduced projection segments for backward. + ct.store(PROJ_OUT, index=(bid_m, 0), tile=pre_accum.astype(PROJ_OUT.dtype)) + ct.store(PROJ_OUT, index=(bid_m, 1), tile=post_accum.astype(PROJ_OUT.dtype)) + + # 2. Compute r = norm / sqrt(K) + denom = ct.full((TILE_SIZE_M, 1), K * 1.0, dtype=ct.float32) + r_val = ct.sqrt(ct.truediv(r_accum, denom)) + + ct.store(R, index=(bid_m, 0), tile=r_val.astype(R.dtype)) + + # 3. Apply compute_h directly into split outputs. + inv_r_eps = 1.0 / (r_val + eps) + bias_pre = ct.load(Bias, index=(0, 0), shape=(1, n), padding_mode=PAD_ZERO) + bias_post = ct.load(Bias, index=(0, 1), shape=(1, n), padding_mode=PAD_ZERO) + bias_pre = ct.astype(bias_pre, ct.float32) + bias_post = ct.astype(bias_post, ct.float32) + + h_pre_linear = pre_accum * alpha_pre * inv_r_eps + bias_pre + h_post_linear = post_accum * alpha_post * inv_r_eps + bias_post + h_pre = _ct_sigmoid(h_pre_linear) + compute_h_eps + h_post = _ct_sigmoid(h_post_linear) * 2.0 + + ct.store(H_PRE, index=(bid_m, 0), tile=h_pre.astype(H_PRE.dtype)) + ct.store(H_POST, index=(bid_m, 0), tile=h_post.astype(H_POST.dtype)) + + for res_chunk in ct.static_iter(range(n)): + res_accum = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) + for split_idx in ct.static_iter(range(SPLIT_K)): + bid_m_k = bid_m + split_idx * num_bid_m + res_tile = ct.load( + Y_acc, + index=(bid_m_k, 2 + res_chunk), + shape=(TILE_SIZE_M, n), + padding_mode=PAD_ZERO, + ) + res_accum = res_accum + ct.astype(res_tile, ct.float32) + + bias_res = ct.load(Bias, index=(0, 2 + res_chunk), shape=(1, n), padding_mode=PAD_ZERO) + bias_res = ct.astype(bias_res, ct.float32) + h_res = res_accum * alpha_res * inv_r_eps + bias_res + ct.store(PROJ_OUT, index=(bid_m, 2 + res_chunk), tile=res_accum.astype(PROJ_OUT.dtype)) + ct.store(H_RES, index=(bid_m, res_chunk), tile=h_res.astype(H_RES.dtype)) + + @ct.kernel + def _ct_proj_rms_bwd_kernel( + A, + B, + NORM, + DD, + DR, + DA, + DB, + M: int, + N: int, + K: int, + eps: float, + TILE_SIZE_M: ConstInt, + TILE_SIZE_N: ConstInt, + TILE_SIZE_K: ConstInt, + ): + zero_pad = ct.PaddingMode.ZERO + tile_k_id = ct.bid(0) + NUM_M_TILES = ct.cdiv(M, TILE_SIZE_M) + accumulator_db = ct.full((TILE_SIZE_K, TILE_SIZE_N), 0.0, dtype=ct.float32) + for tile_m_id in range(NUM_M_TILES): + accumulator_da = ct.full((TILE_SIZE_M, TILE_SIZE_K), 0.0, dtype=ct.float32) + a_tile = ct.load( + A, + index=(tile_m_id, tile_k_id), + shape=(TILE_SIZE_M, TILE_SIZE_K), + padding_mode=zero_pad, + ) + norm_tile = ct.load( + NORM, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), padding_mode=zero_pad + ) + dr_tile = ct.load( + DR, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), padding_mode=zero_pad + ) + accumulator_da = accumulator_da + _ct_rms_dnorm(a_tile, norm_tile, dr_tile, K, eps) + b_tile = ct.load( + B, index=(0, tile_k_id), shape=(TILE_SIZE_N, TILE_SIZE_K), padding_mode=zero_pad + ) + dd_tile = ct.load( + DD, index=(tile_m_id, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), padding_mode=zero_pad + ) + dd_tile = ct.astype(dd_tile, ct.tfloat32) + accumulator_da = ct.mma(dd_tile, b_tile.astype(ct.tfloat32), acc=accumulator_da) + ct.store(DA, index=(tile_m_id, tile_k_id), tile=accumulator_da.astype(DA.dtype)) + accumulator_db = ct.mma( + a_tile.transpose().astype(ct.tfloat32), dd_tile, acc=accumulator_db + ) + ct.store(DB, index=(0, tile_k_id), tile=accumulator_db.transpose().astype(DB.dtype)) + + @ct.kernel + def _ct_proj_rms_bwd_small_k_kernel( + A, B, NORM, DD, DR, DA, DB, M: int, N: int, K: int, eps: float, TILE_N_SIZE: ConstInt + ): + zero_pad = ct.PaddingMode.ZERO + TILE_DB_SIZE_M = 128 + TILE_DB_SIZE_K = 64 + NUM_M_TILES = ct.cdiv(M, TILE_DB_SIZE_M) + NUM_K_TILES = ct.cdiv(K, TILE_DB_SIZE_K) + if ct.bid(1) == 0: + for tile_id in range(ct.bid(0), NUM_K_TILES, ct.num_blocks(0)): + accumulator_db = ct.full((TILE_DB_SIZE_K, TILE_N_SIZE), 0.0, dtype=ct.float32) + for m_tile in range(NUM_M_TILES): + a_tile = ct.load( + A, + index=(m_tile, tile_id), + shape=(TILE_DB_SIZE_M, TILE_DB_SIZE_K), + padding_mode=zero_pad, + ) + dd_tile = ct.load( + DD, + index=(m_tile, 0), + shape=(TILE_DB_SIZE_M, TILE_N_SIZE), + padding_mode=zero_pad, + ) + accumulator_db = ct.mma( + a_tile.transpose().astype(ct.tfloat32), + dd_tile.astype(ct.tfloat32), + acc=accumulator_db, + ) + ct.store( + DB, + index=(0, tile_id), + tile=accumulator_db.transpose().astype(DB.dtype), + allow_tma=False, + ) + TILE_DA_SIZE_M = 128 + TILE_DA_SIZE_K = 256 + NUM_DA_TILES = ct.cdiv(M, TILE_DA_SIZE_M) * ct.cdiv(K, TILE_DA_SIZE_K) + NUM_DA_K_TILES = ct.cdiv(K, TILE_DA_SIZE_K) + if ct.bid(1) == 1: + for tile_id in range(ct.bid(0), NUM_DA_TILES, ct.num_blocks(0)): + b_tile_idx = tile_id % NUM_DA_K_TILES + dd_tile_idx = tile_id // NUM_DA_K_TILES + accumulator_da = ct.full((TILE_DA_SIZE_M, TILE_DA_SIZE_K), 0.0, dtype=ct.float32) + a_tile = ct.load( + A, + index=(dd_tile_idx, b_tile_idx), + shape=(TILE_DA_SIZE_M, TILE_DA_SIZE_K), + padding_mode=zero_pad, + ) + norm_tile = ct.load( + NORM, index=(dd_tile_idx, 0), shape=(TILE_DA_SIZE_M, 1), padding_mode=zero_pad + ) + dr_tile = ct.load( + DR, index=(dd_tile_idx, 0), shape=(TILE_DA_SIZE_M, 1), padding_mode=zero_pad + ) + accumulator_da = accumulator_da + _ct_rms_dnorm( + a_tile.astype(ct.float32), norm_tile, dr_tile, K, eps + ) + b_tile = ct.load( + B, + index=(0, b_tile_idx), + shape=(TILE_N_SIZE, TILE_DA_SIZE_K), + padding_mode=zero_pad, + ) + dd_tile = ct.load( + DD, + index=(dd_tile_idx, 0), + shape=(TILE_DA_SIZE_M, TILE_N_SIZE), + padding_mode=zero_pad, + ) + accumulator_da = ct.mma( + dd_tile.astype(ct.tfloat32), b_tile.astype(ct.tfloat32), acc=accumulator_da + ) + ct.store(DA, index=(dd_tile_idx, b_tile_idx), tile=accumulator_da.astype(DA.dtype)) + + def _next_power_of_2(n: int) -> int: + n -= 1 + n |= n >> 1 + n |= n >> 2 + n |= n >> 4 + n |= n >> 8 + n |= n >> 16 + n |= n >> 32 + n += 1 + return n + + def _proj_rms_fwd_autotune_configs(N): + """Generate autotune search space for proj_rms forward kernel.""" + TILE_N = _next_power_of_2(N) + tile_ms = (32, 64, 128) + tile_ks = (32, 64, 128) + split_ks = (1, 2, 4, 8, 16) + for tile_m in tile_ms: + for tile_k in tile_ks: + for split_k in split_ks: + yield {"TILE_M": tile_m, "TILE_N": TILE_N, "TILE_K": tile_k, 'SPLIT_K': split_k} + + def _default_tile_m(M: int) -> int: + """Pick a tile size that avoids unmasked stores past the M dimension.""" + for tile_m in (128, 64, 32, 16, 8, 4, 2, 1): + if tile_m <= M and M % tile_m == 0: + return tile_m + return 1 + + def _default_proj_rms_fwd_config(M: int, K: int, TILE_N: int): + """Static fallback for skinny MHC projection when autotune cache is absent.""" + split_k = 16 if K >= 16384 else 8 if K >= 8192 else 1 + return _default_tile_m(M), TILE_N, min(128, K), split_k + + # Cache the best config across calls (keyed by M, N, K). + _proj_rms_fwd_best_cfg: dict = {} + + def _cutile_proj_rms_fwd( + x: Tensor, weight: Tensor, eps: float = 1e-6 + ) -> Tuple[Tensor, Tensor, Tensor]: + M, K = x.shape + N = weight.shape[0] + TILE_N = _next_power_of_2(N) + dev = x.device + stream = torch.cuda.current_stream() + + cache_key = (M, N, K) + cached = _proj_rms_fwd_best_cfg.get(cache_key) + + if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + # Use cached best config, or fall back to default if no experimental. + if cached is not None: + tm, tn, tk, split_k = cached + else: + tm, tn, tk, split_k = _default_proj_rms_fwd_config(M, K, TILE_N) + + proj = torch.empty(split_k * M, N, dtype=x.dtype, device=dev) + norm = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + # _ct_proj_rms_fwd_kernel keeps R in its signature; r is computed + # below from the reduced norm. + r = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + + ct.launch( + stream, + (math.ceil(M / tm), split_k), + _ct_proj_rms_fwd_kernel, + (x, weight, proj, norm, r, M, N, K, eps, tm, tn, tk, split_k), + ) + proj = proj.view(split_k, M, N).to(torch.float32).sum(dim=0).to(dtype=x.dtype) + norm = norm.view(split_k, M, 1).to(torch.float32).sum(dim=0).to(dtype=x.dtype) + else: + # Autotune on first call for this shape. + from types import SimpleNamespace + + configs = [SimpleNamespace(**c) for c in _proj_rms_fwd_autotune_configs(N)] + # filter out configs with TILE_K > K or TILE_M > M + configs = [cfg for cfg in configs if cfg.TILE_K <= K and M % cfg.TILE_M == 0] + if len(configs) == 0: + tm, tn, tk, split_k = _default_proj_rms_fwd_config(M, K, TILE_N) + proj = torch.empty(split_k * M, N, dtype=x.dtype, device=dev) + norm = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + # Signature placeholder for the cuTile kernel; not read. + r = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + ct.launch( + stream, + (math.ceil(M / tm), split_k), + _ct_proj_rms_fwd_kernel, + (x, weight, proj, norm, r, M, N, K, eps, tm, tn, tk, split_k), + ) + proj = proj.view(split_k, M, N).to(torch.float32).sum(dim=0).to(dtype=x.dtype) + norm = norm.view(split_k, M, 1).to(torch.float32).sum(dim=0).to(dtype=x.dtype) + else: + mx_split_k = max(cfg.SPLIT_K for cfg in configs) + proj = torch.empty(mx_split_k * M, N, dtype=x.dtype, device=dev) + norm = torch.empty(mx_split_k * M, 1, dtype=x.dtype, device=dev) + # Signature placeholder for autotune launches; not read. + r = torch.empty(mx_split_k * M, 1, dtype=x.dtype, device=dev) + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(M / cfg.TILE_M), cfg.SPLIT_K), + kernel=_ct_proj_rms_fwd_kernel, + args_fn=lambda cfg: ( + x, + weight, + proj, + norm, + r, + M, + N, + K, + eps, + cfg.TILE_M, + cfg.TILE_N, + cfg.TILE_K, + cfg.SPLIT_K, + ), + search_space=configs, + ) + best = tuned.tuned_config + _proj_rms_fwd_best_cfg[cache_key] = ( + best.TILE_M, + best.TILE_N, + best.TILE_K, + best.SPLIT_K, + ) + proj = torch.empty(best.SPLIT_K * M, N, dtype=x.dtype, device=dev) + norm = torch.empty(best.SPLIT_K * M, 1, dtype=x.dtype, device=dev) + # Signature placeholder for the cuTile kernel; not read. + r = torch.empty(best.SPLIT_K * M, 1, dtype=x.dtype, device=dev) + # Re-launch with best config for correct output. + ct.launch( + stream, + (math.ceil(M / best.TILE_M), best.SPLIT_K), + _ct_proj_rms_fwd_kernel, + ( + x, + weight, + proj, + norm, + r, + M, + N, + K, + eps, + best.TILE_M, + best.TILE_N, + best.TILE_K, + best.SPLIT_K, + ), + ) + + proj = proj.view(best.SPLIT_K, M, N).to(torch.float32).sum(dim=0).to(dtype=x.dtype) + norm = norm.view(best.SPLIT_K, M, 1).to(torch.float32).sum(dim=0).to(dtype=x.dtype) + norm = torch.sqrt(norm) + r = 1.0 / (norm / math.sqrt(K) + eps) + return proj, norm, r + + # -- Reduce + compute_h launcher ------------------------------------------ + + def _reduce_compute_h_autotune_configs(M): + """Generate autotune search space for reduce_compute_h kernel.""" + min_tile_m = 16 if M >= 16 else 1 + for tile_m in (128, 64, 32, 16, 8, 4, 2, 1): + if tile_m < min_tile_m: + continue + if tile_m <= M and M % tile_m == 0: + yield tile_m + + def _default_reduce_compute_h_tile_m(M: int) -> int: + """Pick a reduce tile size with enough blocks to cover the GPU.""" + try: + num_sms = torch.cuda.get_device_properties("cuda").multi_processor_count + except Exception: + num_sms = 128 + + valid = [tm for tm in _reduce_compute_h_autotune_configs(M)] + for tm in valid: + if math.ceil(M / tm) >= num_sms: + return tm + return valid[-1] if valid else 1 + + _reduce_compute_h_best_cfg: dict = {} + + def _cutile_reduce_compute_h( + proj_acc: Tensor, + norm_acc: Tensor, + bias: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + n: int, + M: int, + N: int, + K: int, + eps: float, + compute_h_eps: float, + _proj_tile_m: int, + tile_n: int, + split_k: int, + ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: + """Launch reduce split-K + compute_h kernel. + + Returns: + h_pre: [M, n] sigmoid-activated pre weights + h_post: [M, n] 2*sigmoid-activated post weights + h_res: [M, n*n] residual logits + r: [M, 1] r = norm / sqrt(K) + proj_reduced: [M, N] reduced projection (for backward) + """ + dev = proj_acc.device + stream = torch.cuda.current_stream() + + bias_2d = bias.unsqueeze(0).contiguous() # [1, N] + + h_pre_out = torch.empty(M, n, dtype=proj_acc.dtype, device=dev) + h_post_out = torch.empty(M, n, dtype=proj_acc.dtype, device=dev) + h_res_out = torch.empty(M, N - 2 * n, dtype=proj_acc.dtype, device=dev) + r_out = torch.empty(M, 1, dtype=proj_acc.dtype, device=dev) + proj_out = torch.empty(M, N, dtype=proj_acc.dtype, device=dev) + + default_tm = _default_reduce_compute_h_tile_m(M) + cache_key = (M, N, K, n, split_k) + cached = _reduce_compute_h_best_cfg.get(cache_key) + + def _make_args(tm): + return ( + proj_acc, + norm_acc, + bias_2d, + alpha_pre, + alpha_post, + alpha_res, + h_pre_out, + h_post_out, + h_res_out, + r_out, + proj_out, + M, + N, + K, + n, + eps, + compute_h_eps, + tm, + tile_n, + split_k, + ) + + if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + tm = cached if cached is not None else default_tm + ct.launch(stream, (math.ceil(M / tm),), _ct_reduce_compute_h_kernel, _make_args(tm)) + else: + from types import SimpleNamespace + + configs = [SimpleNamespace(TILE_M=tm) for tm in _reduce_compute_h_autotune_configs(M)] + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(M / cfg.TILE_M),), + kernel=_ct_reduce_compute_h_kernel, + args_fn=lambda cfg: _make_args(cfg.TILE_M), + search_space=configs, + ) + best_tm = tuned.tuned_config.TILE_M + _reduce_compute_h_best_cfg[cache_key] = best_tm + ct.launch( + stream, (math.ceil(M / best_tm),), _ct_reduce_compute_h_kernel, _make_args(best_tm) + ) + + return h_pre_out, h_post_out, h_res_out, r_out, proj_out + + # -- Combined proj_rms + compute_h forward -------------------------------- + + def _cutile_proj_rms_compute_h_fwd( + x: Tensor, + weight: Tensor, + bias: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + n: int, + eps: float, + compute_h_eps: float, + ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: + """Fused proj_rms + compute_h forward. + + Launches the existing _ct_proj_rms_fwd_kernel (split-K matmul + partial norm), + then _ct_reduce_compute_h_kernel (reduce + r + activations). + + Returns: + h_pre: [M, n] activated pre weights + h_post: [M, n] activated post weights + h_res: [M, n*n] residual logits + r: [M, 1] r = norm / sqrt(K) + proj_reduced: [M, N] reduced projection (for backward) + """ + M, K = x.shape + N = weight.shape[0] + TILE_N = _next_power_of_2(N) + dev = x.device + stream = torch.cuda.current_stream() + + cache_key = (M, N, K) + cached = _proj_rms_fwd_best_cfg.get(cache_key) + + if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + if cached is not None: + tm, tn, tk, split_k = cached + else: + tm, tn, tk, split_k = _default_proj_rms_fwd_config(M, K, TILE_N) + + proj_acc = torch.empty(split_k * M, N, dtype=x.dtype, device=dev) + norm_acc = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + # _ct_proj_rms_fwd_kernel keeps R in its signature; reduce_compute_h + # computes r from norm_acc. + r_placeholder = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + + ct.launch( + stream, + (math.ceil(M / tm), split_k), + _ct_proj_rms_fwd_kernel, + (x, weight, proj_acc, norm_acc, r_placeholder, M, N, K, eps, tm, tn, tk, split_k), + ) + else: + from types import SimpleNamespace + + configs = [SimpleNamespace(**c) for c in _proj_rms_fwd_autotune_configs(N)] + configs = [cfg for cfg in configs if cfg.TILE_K <= K and M % cfg.TILE_M == 0] + if len(configs) == 0: + tm, tn, tk, split_k = _default_proj_rms_fwd_config(M, K, TILE_N) + proj_acc = torch.empty(split_k * M, N, dtype=x.dtype, device=dev) + norm_acc = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + # Signature placeholder for the cuTile kernel; not read. + r_placeholder = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + ct.launch( + stream, + (math.ceil(M / tm), split_k), + _ct_proj_rms_fwd_kernel, + ( + x, + weight, + proj_acc, + norm_acc, + r_placeholder, + M, + N, + K, + eps, + tm, + tn, + tk, + split_k, + ), + ) + else: + mx_split_k = max(cfg.SPLIT_K for cfg in configs) + proj_acc = torch.empty(mx_split_k * M, N, dtype=x.dtype, device=dev) + norm_acc = torch.empty(mx_split_k * M, 1, dtype=x.dtype, device=dev) + # Signature placeholder for autotune launches; not read. + r_placeholder = torch.empty(mx_split_k * M, 1, dtype=x.dtype, device=dev) + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(M / cfg.TILE_M), cfg.SPLIT_K), + kernel=_ct_proj_rms_fwd_kernel, + args_fn=lambda cfg: ( + x, + weight, + proj_acc, + norm_acc, + r_placeholder, + M, + N, + K, + eps, + cfg.TILE_M, + cfg.TILE_N, + cfg.TILE_K, + cfg.SPLIT_K, + ), + search_space=configs, + ) + best = tuned.tuned_config + _proj_rms_fwd_best_cfg[cache_key] = ( + best.TILE_M, + best.TILE_N, + best.TILE_K, + best.SPLIT_K, + ) + tm, tn, tk, split_k = best.TILE_M, best.TILE_N, best.TILE_K, best.SPLIT_K + + proj_acc = torch.empty(split_k * M, N, dtype=x.dtype, device=dev) + norm_acc = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + # Signature placeholder for the cuTile kernel; not read. + r_placeholder = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + ct.launch( + stream, + (math.ceil(M / tm), split_k), + _ct_proj_rms_fwd_kernel, + ( + x, + weight, + proj_acc, + norm_acc, + r_placeholder, + M, + N, + K, + eps, + tm, + tn, + tk, + split_k, + ), + ) + + # Launch reduce + compute_h kernel + h_pre, h_post, h_res, r, proj_reduced = _cutile_reduce_compute_h( + proj_acc, + norm_acc, + bias, + alpha_pre, + alpha_post, + alpha_res, + n, + M, + N, + K, + eps, + compute_h_eps, + tm, + TILE_N, + split_k, + ) + return h_pre, h_post, h_res, r, proj_reduced + + def _proj_rms_bwd_autotune_configs(N): + """Generate autotune search space for proj_rms backward kernel (K >= 8192 path).""" + TILE_N = _next_power_of_2(N) + tile_ms = (32, 64, 128) + tile_ks = (32, 64, 128, 256) + for tile_m in tile_ms: + for tile_k in tile_ks: + yield {"TILE_SIZE_M": tile_m, "TILE_SIZE_N": TILE_N, "TILE_SIZE_K": tile_k} + + _proj_rms_bwd_best_cfg: dict = {} + + def _cutile_proj_rms_bwd( + grad_proj: Tensor, + grad_r: Tensor, + x: Tensor, + weight: Tensor, + norm: Tensor, + eps: float = 1e-6, + ) -> Tuple[Tensor, Tensor]: + M, K = x.shape + N = weight.shape[0] + da = torch.empty_like(x) + db = torch.empty_like(weight) + TILE_SIZE_N = _next_power_of_2(N) + assert TILE_SIZE_N <= 256, f"TILE_SIZE_N too large: {TILE_SIZE_N}" + stream = torch.cuda.current_stream() + + if K >= 8192: + cache_key = (M, N, K) + cached = _proj_rms_bwd_best_cfg.get(cache_key) + + if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + if cached is not None: + tm, tn, tk = cached + else: + tm, tn, tk = 128, TILE_SIZE_N, 128 + ct.launch( + stream, + (math.ceil(K / tk), 1), + _ct_proj_rms_bwd_kernel, + (x, weight, norm, grad_proj, grad_r, da, db, M, N, K, eps, tm, tn, tk), + ) + else: + from types import SimpleNamespace + + configs = [SimpleNamespace(**c) for c in _proj_rms_bwd_autotune_configs(N)] + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(K / cfg.TILE_SIZE_K), 1), + kernel=_ct_proj_rms_bwd_kernel, + args_fn=lambda cfg: ( + x, + weight, + norm, + grad_proj, + grad_r, + da, + db, + M, + N, + K, + eps, + cfg.TILE_SIZE_M, + cfg.TILE_SIZE_N, + cfg.TILE_SIZE_K, + ), + search_space=configs, + ) + best = tuned.tuned_config + _proj_rms_bwd_best_cfg[cache_key] = ( + best.TILE_SIZE_M, + best.TILE_SIZE_N, + best.TILE_SIZE_K, + ) + ct.launch( + stream, + (math.ceil(K / best.TILE_SIZE_K), 1), + _ct_proj_rms_bwd_kernel, + ( + x, + weight, + norm, + grad_proj, + grad_r, + da, + db, + M, + N, + K, + eps, + best.TILE_SIZE_M, + best.TILE_SIZE_N, + best.TILE_SIZE_K, + ), + ) + else: + num_sms = torch.cuda.get_device_properties("cuda").multi_processor_count + grid = (num_sms, 2, 1) + ct.launch( + stream, + grid, + _ct_proj_rms_bwd_small_k_kernel, + (x, weight, norm, grad_proj, grad_r, da, db, M, N, K, eps, TILE_SIZE_N), + ) + return da, db + + # -- Fused compute_h + proj_rms backward kernels ---------------------------- + + @ct.kernel + def _ct_fused_grad_h_proj_kernel( + GRAD_H_PRE, # [M, n] + GRAD_H_POST, # [M, n] + GRAD_H_RES, # [M, n*n] + H_PRE, # [M, n] + H_POST, # [M, n] + PROJ, # [M, N] + R, # [M, 1] + GRAD_R_EXT, # [M, 1] + Alpha_pre, # [1] + Alpha_post, # [1] + Alpha_res, # [1] + GRAD_H, # [M, TILE_SIZE_N] output + GRAD_PROJ, # [M, TILE_SIZE_N] output + GRAD_R_TOTAL, # [M, 1] output + M: int, + N: int, + n: ConstInt, + eps: float, + compute_h_eps: float, + TILE_SIZE_M: ConstInt, + TILE_SIZE_N: ConstInt, + HAS_GRAD_H_PRE: ConstInt, + HAS_GRAD_H_POST: ConstInt, + HAS_GRAD_H_RES: ConstInt, + HAS_GRAD_R_EXT: ConstInt, + ): + """Precompute grad_h, grad_proj, and grad_r_total for downstream backward kernels. + + Grid: (ceil(M / TILE_SIZE_M),). + """ + tile_m_id = ct.bid(0) + + alpha_pre = ct.load(Alpha_pre, index=(0,), shape=(1,)).item() + alpha_post = ct.load(Alpha_post, index=(0,), shape=(1,)).item() + alpha_res = ct.load(Alpha_res, index=(0,), shape=(1,)).item() + + r_tile = ct.load(R, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), padding_mode=PAD_ZERO) + r_tile = ct.astype(r_tile, ct.float32) + + r_eps = r_tile + eps + inv_r_eps = 1.0 / r_eps + grad_r_from_h = ct.full((TILE_SIZE_M, 1), 0.0, dtype=ct.float32) + + # Clear the padded columns once inside this kernel. Valid columns are + # overwritten by the segment stores below. + zero_full = ct.full((TILE_SIZE_M, TILE_SIZE_N), 0.0, dtype=ct.float32) + ct.store(GRAD_H, index=(tile_m_id, 0), tile=zero_full.astype(GRAD_H.dtype)) + ct.store(GRAD_PROJ, index=(tile_m_id, 0), tile=zero_full.astype(GRAD_PROJ.dtype)) + + if HAS_GRAD_H_PRE: + gy_pre = ct.load( + GRAD_H_PRE, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO + ) + gy_pre = ct.astype(gy_pre, ct.float32) + else: + gy_pre = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) + h_pre = ct.load(H_PRE, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO) + h_pre = ct.astype(h_pre, ct.float32) + proj_pre = ct.load( + PROJ, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO + ) + proj_pre = ct.astype(proj_pre, ct.float32) + sigmoid_pre = h_pre - compute_h_eps + grad_h_pre = gy_pre * sigmoid_pre * (1.0 - sigmoid_pre) + grad_proj_pre = grad_h_pre * alpha_pre * inv_r_eps + grad_r_from_h += ct.sum( + grad_h_pre * proj_pre * alpha_pre * (-inv_r_eps * inv_r_eps), axis=1, keepdims=True + ) + ct.store(GRAD_H, index=(tile_m_id, 0), tile=grad_h_pre.astype(GRAD_H.dtype)) + ct.store(GRAD_PROJ, index=(tile_m_id, 0), tile=grad_proj_pre.astype(GRAD_PROJ.dtype)) + + if HAS_GRAD_H_POST: + gy_post = ct.load( + GRAD_H_POST, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO + ) + gy_post = ct.astype(gy_post, ct.float32) + else: + gy_post = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) + h_post = ct.load( + H_POST, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO + ) + h_post = ct.astype(h_post, ct.float32) + proj_post = ct.load( + PROJ, index=(tile_m_id, 1), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO + ) + proj_post = ct.astype(proj_post, ct.float32) + sigmoid_post = h_post * 0.5 + grad_h_post = gy_post * sigmoid_post * (1.0 - sigmoid_post) * 2.0 + grad_proj_post = grad_h_post * alpha_post * inv_r_eps + grad_r_from_h += ct.sum( + grad_h_post * proj_post * alpha_post * (-inv_r_eps * inv_r_eps), axis=1, keepdims=True + ) + ct.store(GRAD_H, index=(tile_m_id, 1), tile=grad_h_post.astype(GRAD_H.dtype)) + ct.store(GRAD_PROJ, index=(tile_m_id, 1), tile=grad_proj_post.astype(GRAD_PROJ.dtype)) + + for res_chunk in ct.static_iter(range(n)): + if HAS_GRAD_H_RES: + grad_h_res = ct.load( + GRAD_H_RES, + index=(tile_m_id, res_chunk), + shape=(TILE_SIZE_M, n), + padding_mode=PAD_ZERO, + ) + grad_h_res = ct.astype(grad_h_res, ct.float32) + else: + grad_h_res = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) + proj_res = ct.load( + PROJ, + index=(tile_m_id, 2 + res_chunk), + shape=(TILE_SIZE_M, n), + padding_mode=PAD_ZERO, + ) + proj_res = ct.astype(proj_res, ct.float32) + grad_proj_res = grad_h_res * alpha_res * inv_r_eps + grad_r_from_h += ct.sum( + grad_h_res * proj_res * alpha_res * (-inv_r_eps * inv_r_eps), axis=1, keepdims=True + ) + ct.store(GRAD_H, index=(tile_m_id, 2 + res_chunk), tile=grad_h_res.astype(GRAD_H.dtype)) + ct.store( + GRAD_PROJ, + index=(tile_m_id, 2 + res_chunk), + tile=grad_proj_res.astype(GRAD_PROJ.dtype), + ) + + if HAS_GRAD_R_EXT: + grad_r_ext_tile = ct.load( + GRAD_R_EXT, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), padding_mode=PAD_ZERO + ) + grad_r_ext_tile = ct.astype(grad_r_ext_tile, ct.float32) + else: + grad_r_ext_tile = ct.full((TILE_SIZE_M, 1), 0.0, dtype=ct.float32) + grad_r_total = grad_r_from_h + grad_r_ext_tile + + ct.store(GRAD_R_TOTAL, index=(tile_m_id, 0), tile=grad_r_total.astype(GRAD_R_TOTAL.dtype)) + + @ct.kernel + def _ct_fused_grad_x_weight_kernel( + X, # [M, K] + WEIGHT, # [N, K] + GRAD_PROJ, # [M, TILE_SIZE_N] precomputed + GRAD_R_TOTAL, # [M, 1] precomputed + R, # [M, 1] + GRAD_X, # [M, K] output + GRAD_WEIGHT, # [N, K] output + M: int, + N: int, + K: int, + TILE_SIZE_M: ConstInt, + TILE_SIZE_N: ConstInt, + TILE_SIZE_K: ConstInt, + ): + """Compute grad_x and grad_weight simultaneously. + + Grid: (ceil(K / TILE_SIZE_K),). + Each block handles one K-tile and loops over all M-tiles. + Per M-tile: computes and stores grad_x, accumulates grad_weight. + """ + tile_k_id = ct.bid(0) + NUM_M_TILES = ct.cdiv(M, TILE_SIZE_M) + + # Load weight tile once — only depends on K-tile + weight_tile = ct.load( + WEIGHT, index=(0, tile_k_id), shape=(TILE_SIZE_N, TILE_SIZE_K), padding_mode=PAD_ZERO + ) + + acc_grad_weight = ct.full((TILE_SIZE_K, TILE_SIZE_N), 0.0, dtype=ct.float32) + + for tile_m_id in range(NUM_M_TILES): + grad_proj_tile = ct.load( + GRAD_PROJ, + index=(tile_m_id, 0), + shape=(TILE_SIZE_M, TILE_SIZE_N), + padding_mode=PAD_ZERO, + ) + x_tile = ct.load( + X, + index=(tile_m_id, tile_k_id), + shape=(TILE_SIZE_M, TILE_SIZE_K), + padding_mode=PAD_ZERO, + ) + grad_r_total = ct.load( + GRAD_R_TOTAL, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), padding_mode=PAD_ZERO + ) + r_tile = ct.load(R, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), padding_mode=PAD_ZERO) + r_tile = ct.astype(r_tile, ct.float32) + + # grad_x = grad_proj @ weight + grad_r_total * x / (r * K) + inv_rK = 1.0 / (r_tile * K) + acc_grad_x = (grad_r_total * inv_rK) * ct.astype(x_tile, ct.float32) + acc_grad_x = ct.mma( + grad_proj_tile.astype(ct.tfloat32), weight_tile.astype(ct.tfloat32), acc=acc_grad_x + ) + ct.store(GRAD_X, index=(tile_m_id, tile_k_id), tile=acc_grad_x.astype(GRAD_X.dtype)) + + # Accumulate grad_weight += x.T @ grad_proj + acc_grad_weight = ct.mma( + x_tile.transpose().astype(ct.tfloat32), + grad_proj_tile.astype(ct.tfloat32), + acc=acc_grad_weight, + ) + + ct.store( + GRAD_WEIGHT, + index=(0, tile_k_id), + tile=acc_grad_weight.transpose().astype(GRAD_WEIGHT.dtype), + ) + + @ct.kernel + def _ct_scalar_grads_partials_kernel( + GRAD_H, # [M, TILE_SIZE_N] precomputed + PROJ, # [M, N] + R, # [M, 1] + GRAD_ALPHA_PRE_PARTIALS, # [num_m_blocks, 1] output + GRAD_ALPHA_POST_PARTIALS, # [num_m_blocks, 1] output + GRAD_ALPHA_RES_PARTIALS, # [num_m_blocks, 1] output + GRAD_BIAS_PARTIALS, # [num_m_blocks, TILE_SIZE_N] output + M: int, + N: int, + n: int, + eps: float, + TILE_SIZE_M: ConstInt, + TILE_SIZE_N: ConstInt, + ): + """Compute per-M-tile scalar-gradient partials. + + Grid: (ceil(M / TILE_SIZE_M),). Each block processes one M-tile. + """ + bid_m = ct.bid(0) + + offsets = ct.arange(TILE_SIZE_N, dtype=ct.int32) + one = ct.full((TILE_SIZE_N,), 1.0, dtype=ct.float32) + zero = ct.full((TILE_SIZE_N,), 0.0, dtype=ct.float32) + mask_pre = ct.where(ct.less(offsets, n), one, zero) + mask_post = ct.where(ct.less(offsets, 2 * n), one, zero) - mask_pre + mask_res = one - mask_pre - mask_post + + mask_pre_2d = ct.reshape(mask_pre, (1, TILE_SIZE_N)) + mask_post_2d = ct.reshape(mask_post, (1, TILE_SIZE_N)) + mask_res_2d = ct.reshape(mask_res, (1, TILE_SIZE_N)) + + grad_h = ct.load( + GRAD_H, index=(bid_m, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), padding_mode=PAD_ZERO + ) + proj_tile = ct.load( + PROJ, index=(bid_m, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), padding_mode=PAD_ZERO + ) + proj_tile = ct.astype(proj_tile, ct.float32) + r_tile = ct.load(R, index=(bid_m, 0), shape=(TILE_SIZE_M, 1), padding_mode=PAD_ZERO) + r_tile = ct.astype(r_tile, ct.float32) + + r_eps = r_tile + eps + inv_r_eps = 1.0 / r_eps + + ga_all = grad_h * proj_tile * inv_r_eps + ga_pre = ct.reshape(ct.sum(ga_all * mask_pre_2d), (1, 1)) + ga_post = ct.reshape(ct.sum(ga_all * mask_post_2d), (1, 1)) + ga_res = ct.reshape(ct.sum(ga_all * mask_res_2d), (1, 1)) + partial_gb = ct.sum(grad_h, axis=0, keepdims=False) + ct.store( + GRAD_ALPHA_PRE_PARTIALS, + index=(bid_m, 0), + tile=ga_pre.astype(GRAD_ALPHA_PRE_PARTIALS.dtype), + ) + ct.store( + GRAD_ALPHA_POST_PARTIALS, + index=(bid_m, 0), + tile=ga_post.astype(GRAD_ALPHA_POST_PARTIALS.dtype), + ) + ct.store( + GRAD_ALPHA_RES_PARTIALS, + index=(bid_m, 0), + tile=ga_res.astype(GRAD_ALPHA_RES_PARTIALS.dtype), + ) + ct.store( + GRAD_BIAS_PARTIALS, + index=(bid_m, 0), + tile=ct.reshape(partial_gb, (1, TILE_SIZE_N)).astype(GRAD_BIAS_PARTIALS.dtype), + ) + + @ct.kernel + def _ct_scalar_grads_reduce_kernel( + GRAD_ALPHA_PRE_PARTIALS, # [num_m_blocks, 1] + GRAD_ALPHA_POST_PARTIALS, # [num_m_blocks, 1] + GRAD_ALPHA_RES_PARTIALS, # [num_m_blocks, 1] + GRAD_BIAS_PARTIALS, # [num_m_blocks, TILE_SIZE_N] + GRAD_ALPHA_PRE, # [1, 1] output + GRAD_ALPHA_POST, # [1, 1] output + GRAD_ALPHA_RES, # [1, 1] output + GRAD_BIAS, # [1, TILE_SIZE_N] output + NUM_M_BLOCKS: int, + TILE_SIZE_N: ConstInt, + ): + """Reduce scalar-gradient partials and write final dtype outputs.""" + acc_pre = ct.full((1, 1), 0.0, dtype=ct.float32) + acc_post = ct.full((1, 1), 0.0, dtype=ct.float32) + acc_res = ct.full((1, 1), 0.0, dtype=ct.float32) + acc_bias = ct.full((1, TILE_SIZE_N), 0.0, dtype=ct.float32) + + for bid_m in range(NUM_M_BLOCKS): + acc_pre += ct.load( + GRAD_ALPHA_PRE_PARTIALS, index=(bid_m, 0), shape=(1, 1), padding_mode=PAD_ZERO + ).astype(ct.float32) + acc_post += ct.load( + GRAD_ALPHA_POST_PARTIALS, index=(bid_m, 0), shape=(1, 1), padding_mode=PAD_ZERO + ).astype(ct.float32) + acc_res += ct.load( + GRAD_ALPHA_RES_PARTIALS, index=(bid_m, 0), shape=(1, 1), padding_mode=PAD_ZERO + ).astype(ct.float32) + acc_bias += ct.load( + GRAD_BIAS_PARTIALS, index=(bid_m, 0), shape=(1, TILE_SIZE_N), padding_mode=PAD_ZERO + ).astype(ct.float32) + + ct.store(GRAD_ALPHA_PRE, index=(0, 0), tile=acc_pre.astype(GRAD_ALPHA_PRE.dtype)) + ct.store(GRAD_ALPHA_POST, index=(0, 0), tile=acc_post.astype(GRAD_ALPHA_POST.dtype)) + ct.store(GRAD_ALPHA_RES, index=(0, 0), tile=acc_res.astype(GRAD_ALPHA_RES.dtype)) + ct.store(GRAD_BIAS, index=(0, 0), tile=acc_bias.astype(GRAD_BIAS.dtype)) + + @ct.kernel + def _ct_fused_compute_h_proj_rms_bwd_small_k_kernel( + X, # [M, K] + WEIGHT, # [N, K] + GRAD_PROJ, # [M, TILE_N] precomputed + GRAD_R_TOTAL, # [M, 1] precomputed + R, # [M, 1] + GRAD_X, # [M, K] output + GRAD_WEIGHT, # [N, K] output + M: int, + N: int, + K: int, + TILE_N_SIZE: ConstInt, + ): + """Fused backward (small K path) with work-stealing. + + Grid: (num_sms, 2). + bid(1)==0: grad_weight via work-stealing over K-tiles, loops M. + bid(1)==1: grad_x via work-stealing over (M×K) tiles. + Scalar gradients are computed by the separate partial/reduce kernels. + """ + zero_pad = ct.PaddingMode.ZERO + + TILE_DB_SIZE_M = 128 + TILE_DB_SIZE_K = 64 + NUM_M_TILES = ct.cdiv(M, TILE_DB_SIZE_M) + NUM_K_TILES = ct.cdiv(K, TILE_DB_SIZE_K) + + if ct.bid(1) == 0: + # --- grad_weight path --- + for tile_id in range(ct.bid(0), NUM_K_TILES, ct.num_blocks(0)): + accumulator_db = ct.full((TILE_DB_SIZE_K, TILE_N_SIZE), 0.0, dtype=ct.float32) + for m_tile in range(NUM_M_TILES): + x_tile = ct.load( + X, + index=(m_tile, tile_id), + shape=(TILE_DB_SIZE_M, TILE_DB_SIZE_K), + padding_mode=zero_pad, + ) + grad_proj_tile = ct.load( + GRAD_PROJ, + index=(m_tile, 0), + shape=(TILE_DB_SIZE_M, TILE_N_SIZE), + padding_mode=zero_pad, + ) + + accumulator_db = ct.mma( + x_tile.transpose().astype(ct.tfloat32), + grad_proj_tile.astype(ct.tfloat32), + acc=accumulator_db, + ) + + ct.store( + GRAD_WEIGHT, + index=(0, tile_id), + tile=accumulator_db.transpose().astype(GRAD_WEIGHT.dtype), + allow_tma=False, + ) + + TILE_DA_SIZE_M = 128 + TILE_DA_SIZE_K = 256 + NUM_DA_TILES = ct.cdiv(M, TILE_DA_SIZE_M) * ct.cdiv(K, TILE_DA_SIZE_K) + NUM_DA_K_TILES = ct.cdiv(K, TILE_DA_SIZE_K) + + if ct.bid(1) == 1: + # --- grad_x path --- + for tile_id in range(ct.bid(0), NUM_DA_TILES, ct.num_blocks(0)): + b_tile_idx = tile_id % NUM_DA_K_TILES + dd_tile_idx = tile_id // NUM_DA_K_TILES + + grad_proj_tile = ct.load( + GRAD_PROJ, + index=(dd_tile_idx, 0), + shape=(TILE_DA_SIZE_M, TILE_N_SIZE), + padding_mode=zero_pad, + ) + grad_r_total = ct.load( + GRAD_R_TOTAL, + index=(dd_tile_idx, 0), + shape=(TILE_DA_SIZE_M, 1), + padding_mode=zero_pad, + ) + r_tile = ct.load( + R, index=(dd_tile_idx, 0), shape=(TILE_DA_SIZE_M, 1), padding_mode=zero_pad + ) + r_tile = ct.astype(r_tile, ct.float32) + + x_tile = ct.load( + X, + index=(dd_tile_idx, b_tile_idx), + shape=(TILE_DA_SIZE_M, TILE_DA_SIZE_K), + padding_mode=zero_pad, + ) + inv_rK = 1.0 / (r_tile * K) + accumulator_da = (grad_r_total * inv_rK) * ct.astype(x_tile, ct.float32) + + weight_tile = ct.load( + WEIGHT, + index=(0, b_tile_idx), + shape=(TILE_N_SIZE, TILE_DA_SIZE_K), + padding_mode=zero_pad, + ) + accumulator_da = ct.mma( + grad_proj_tile.astype(ct.tfloat32), + weight_tile.astype(ct.tfloat32), + acc=accumulator_da, + ) + ct.store( + GRAD_X, + index=(dd_tile_idx, b_tile_idx), + tile=accumulator_da.astype(GRAD_X.dtype), + ) + + def _fused_grad_x_weight_autotune_configs(N): + """Autotune search space for fused grad_x + grad_weight kernel.""" + TILE_N = _next_power_of_2(N) + tile_ms = (32, 64, 128) + tile_ks = (32, 64, 128, 256) + for tile_m in tile_ms: + for tile_k in tile_ks: + yield {"TILE_SIZE_M": tile_m, "TILE_SIZE_N": TILE_N, "TILE_SIZE_K": tile_k} + + _fused_grad_x_weight_best_cfg: dict = {} + + def _cutile_fused_compute_h_proj_rms_bwd( + x: Tensor, + weight: Tensor, + grad_h_pre: Tensor, + grad_h_post: Tensor, + grad_h_res: Tensor, + h_pre: Tensor, + h_post: Tensor, + h_res: Tensor, + proj: Tensor, + r: Tensor, + grad_r_ext: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + bias: Tensor, + n: int, + eps: float, + compute_h_eps: float, + ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + """Fused compute_h + proj_rms backward. + + Returns: + grad_x: [M, K] + grad_weight: [N, K] + grad_alpha_pre: [1] + grad_alpha_post: [1] + grad_alpha_res: [1] + grad_bias: [N] + """ + M, K = x.shape + N = weight.shape[0] + TILE_N = _next_power_of_2(N) + assert TILE_N <= 256, f"TILE_SIZE_N too large: {TILE_N}" + dev = x.device + stream = torch.cuda.current_stream() + + grad_x = torch.empty_like(x) + grad_weight = torch.empty_like(weight) + has_grad_r_ext = grad_r_ext is not None + has_grad_r_ext_flag = int(has_grad_r_ext) + grad_r_ext_arg = grad_r_ext if has_grad_r_ext else r + has_grad_h_pre = grad_h_pre is not None + has_grad_h_post = grad_h_post is not None + has_grad_h_res = grad_h_res is not None + grad_h_pre_arg = grad_h_pre if has_grad_h_pre else h_pre + grad_h_post_arg = grad_h_post if has_grad_h_post else h_post + grad_h_res_arg = grad_h_res if has_grad_h_res else h_res + + # 0. Precompute grad_h, grad_proj, grad_r_total + grad_h_buf = torch.empty(M, TILE_N, dtype=torch.float32, device=dev) + grad_proj_buf = torch.empty(M, TILE_N, dtype=torch.float32, device=dev) + grad_r_total_buf = torch.empty(M, 1, dtype=torch.float32, device=dev) + + tile_m_precomp = _default_tile_m(M) + ct.launch( + stream, + (math.ceil(M / tile_m_precomp),), + _ct_fused_grad_h_proj_kernel, + ( + grad_h_pre_arg, + grad_h_post_arg, + grad_h_res_arg, + h_pre, + h_post, + proj, + r, + grad_r_ext_arg, + alpha_pre, + alpha_post, + alpha_res, + grad_h_buf, + grad_proj_buf, + grad_r_total_buf, + M, + N, + n, + eps, + compute_h_eps, + tile_m_precomp, + TILE_N, + int(has_grad_h_pre), + int(has_grad_h_post), + int(has_grad_h_res), + has_grad_r_ext_flag, + ), + ) + + if K >= 8192: + # 1. Fused grad_x + grad_weight kernel — 1D grid (K-tiles), loops M + cache_key = ('grad_x_weight', M, N, K) + cached = _fused_grad_x_weight_best_cfg.get(cache_key) + + if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + if cached is not None: + tm, tn, tk = cached + else: + tm, tn, tk = 128, TILE_N, 128 + ct.launch( + stream, + (math.ceil(K / tk),), + _ct_fused_grad_x_weight_kernel, + ( + x, + weight, + grad_proj_buf, + grad_r_total_buf, + r, + grad_x, + grad_weight, + M, + N, + K, + tm, + tn, + tk, + ), + ) + else: + from types import SimpleNamespace + + configs = [SimpleNamespace(**c) for c in _fused_grad_x_weight_autotune_configs(N)] + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(K / cfg.TILE_SIZE_K),), + kernel=_ct_fused_grad_x_weight_kernel, + args_fn=lambda cfg: ( + x, + weight, + grad_proj_buf, + grad_r_total_buf, + r, + grad_x, + grad_weight, + M, + N, + K, + cfg.TILE_SIZE_M, + cfg.TILE_SIZE_N, + cfg.TILE_SIZE_K, + ), + search_space=configs, + ) + best = tuned.tuned_config + _fused_grad_x_weight_best_cfg[cache_key] = ( + best.TILE_SIZE_M, + best.TILE_SIZE_N, + best.TILE_SIZE_K, + ) + ct.launch( + stream, + (math.ceil(K / best.TILE_SIZE_K),), + _ct_fused_grad_x_weight_kernel, + ( + x, + weight, + grad_proj_buf, + grad_r_total_buf, + r, + grad_x, + grad_weight, + M, + N, + K, + best.TILE_SIZE_M, + best.TILE_SIZE_N, + best.TILE_SIZE_K, + ), + ) + else: + num_sms = torch.cuda.get_device_properties("cuda").multi_processor_count + ct.launch( + stream, + (num_sms, 2, 1), + _ct_fused_compute_h_proj_rms_bwd_small_k_kernel, + ( + x, + weight, + grad_proj_buf, + grad_r_total_buf, + r, + grad_x, + grad_weight, + M, + N, + K, + TILE_N, + ), + ) + + # 2. Separate lightweight kernel for scalar gradients (grad_alpha, grad_bias) + tile_m_scalar = min(128, M) + num_m_blocks = math.ceil(M / tile_m_scalar) + grad_alpha_pre_partials = torch.empty(num_m_blocks, 1, dtype=torch.float32, device=dev) + grad_alpha_post_partials = torch.empty(num_m_blocks, 1, dtype=torch.float32, device=dev) + grad_alpha_res_partials = torch.empty(num_m_blocks, 1, dtype=torch.float32, device=dev) + grad_bias_partials = torch.empty(num_m_blocks, TILE_N, dtype=torch.float32, device=dev) + grad_alpha_pre = torch.empty(1, 1, dtype=alpha_pre.dtype, device=dev) + grad_alpha_post = torch.empty(1, 1, dtype=alpha_post.dtype, device=dev) + grad_alpha_res = torch.empty(1, 1, dtype=alpha_res.dtype, device=dev) + grad_bias = torch.empty(1, TILE_N, dtype=bias.dtype, device=dev) + + ct.launch( + stream, + (num_m_blocks,), + _ct_scalar_grads_partials_kernel, + ( + grad_h_buf, + proj, + r, + grad_alpha_pre_partials, + grad_alpha_post_partials, + grad_alpha_res_partials, + grad_bias_partials, + M, + N, + n, + eps, + tile_m_scalar, + TILE_N, + ), + ) + ct.launch( + stream, + (1,), + _ct_scalar_grads_reduce_kernel, + ( + grad_alpha_pre_partials, + grad_alpha_post_partials, + grad_alpha_res_partials, + grad_bias_partials, + grad_alpha_pre, + grad_alpha_post, + grad_alpha_res, + grad_bias, + num_m_blocks, + TILE_N, + ), + ) + + return ( + grad_x, + grad_weight, + grad_alpha_pre.view_as(alpha_pre), + grad_alpha_post.view_as(alpha_post), + grad_alpha_res.view_as(alpha_res), + grad_bias.view(-1)[:N], + ) + + +# ============================================================================ +# Unified public dispatch +# ============================================================================ +# The public fused API chooses the fastest validated backend per operation: +# +# sinkhorn fwd/bwd: Triton -> cuTile -> torch +# h_post_bda fwd/bwd: Triton -> cuTile -> torch +# h_aggregate fwd: Triton -> cuTile -> torch +# h_aggregate bwd: cuTile -> torch +# proj_rms/proj_rms_compute_h: cuTile -> torch +# +# Runtime CUDA launch failures are intentionally not swallowed; after such an +# error the CUDA context may not be safely reusable for fallback work. +# ============================================================================ + +from megatron.core.transformer.hyper_connection import ( + native_fused_add_3, + native_h_aggregate, + native_h_post_bda, + native_proj_rms, + native_sinkhorn, +) + +_BACKEND_INFO_LOGGED = False + + +def _select_triton_cutile_native(triton_impl) -> str: + if triton_impl is not None: + return "triton" + if is_cutile_available(): + return "cutile" + return "native" + + +def _mhc_backend_status() -> Tuple[str, bool]: + """Return backend description and whether every backend is native.""" + sinkhorn = _select_triton_cutile_native(_get_triton_sinkhorn()) + h_aggregate_fwd = _select_triton_cutile_native(_get_triton_h_aggregate_fwd()) + h_aggregate_bwd = "cutile" if is_cutile_available() else "native" + h_post_bda_fwd = _select_triton_cutile_native(_get_triton_h_post_bda_fwd()) + h_post_bda_bwd = _select_triton_cutile_native(_get_triton_h_post_bda_bwd()) + proj_rms = "cutile" if is_cutile_available() else "native" + selected = ( + sinkhorn, + h_aggregate_fwd, + h_aggregate_bwd, + h_post_bda_fwd, + h_post_bda_bwd, + proj_rms, + ) + message = ( + f"MHC_FORCE_BACKEND={_MHC_FORCED_BACKEND}; " + f"sinkhorn={sinkhorn}; " + f"h_aggregate=fwd:{h_aggregate_fwd},bwd:{h_aggregate_bwd}; " + f"h_post_bda=fwd:{h_post_bda_fwd},bwd:{h_post_bda_bwd}; " + f"proj_rms={proj_rms}; " + f"proj_rms_compute_h={proj_rms}" + ) + return message, all(backend == "native" for backend in selected) + + +def _mhc_backend_selection() -> str: + """Return a concise description of the selected mHC fused backends.""" + message, _ = _mhc_backend_status() + return message + + +def log_fused_mhc_backend_once() -> None: + """Log the fused mHC backend selection once per process.""" + _raise_mhc_backend_validation_error() + global _BACKEND_INFO_LOGGED + if _BACKEND_INFO_LOGGED: + return + _BACKEND_INFO_LOGGED = True + backend_selection, all_native = _mhc_backend_status() + log_single_rank( + logger, + logging.WARNING if all_native else logging.INFO, + f"[mHC] fused backend selection: {backend_selection}", + ) + if all_native and safe_get_rank() == 0: + warnings.warn( + "[mHC] No accelerated mHC backend is available; falling back to native torch " + "implementations. The fallback is functionally equivalent, but may not provide " + "the performance benefits of fused mHC backends.", + UserWarning, + stacklevel=2, + ) + + +def fused_add_3(a: Tensor, b: Tensor, c: Tensor) -> Tensor: + """Add three tensors using the native torch.compile-backed implementation.""" + return native_fused_add_3(a, b, c) + + +def _get_triton_sinkhorn(): + if not _TRITON_AVAILABLE: + return None + return _TRITON_IMPLS["sinkhorn"] + + +def _get_triton_h_aggregate_fwd(): + if not _TRITON_AVAILABLE: + return None + return _TRITON_IMPLS["h_aggregate_fwd"] + + +def _get_triton_h_post_bda_fwd(): + if not _TRITON_AVAILABLE: + return None + return _TRITON_IMPLS["h_post_bda_fwd"] + + +def _get_triton_h_post_bda_bwd(): + if not _TRITON_AVAILABLE: + return None + return _TRITON_IMPLS["h_post_bda_bwd"] + + +def _torch_h_aggregate_bwd(grad_output: Tensor, x: Tensor, h_pre: Tensor) -> Tuple[Tensor, Tensor]: + grad_output_expanded = grad_output.unsqueeze(2) + grad_x = grad_output_expanded * h_pre.unsqueeze(-1) + grad_h = torch.sum(grad_output_expanded * x, dim=-1) + return grad_x.to(dtype=x.dtype), grad_h.to(dtype=h_pre.dtype) + + +@torch.compile +def _torch_h_post_bda_bwd( + grad_output: Tensor, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + x: Tensor, + bias: Optional[Tensor], +) -> Tuple[Tensor, Tensor, Tensor, Tensor, Optional[Tensor]]: + s, b, n, C = original_residual.shape + sb = s * b + go = grad_output.reshape(sb, n, C) + hr = h_res.reshape(sb, n, n) + orig = original_residual.reshape(sb, n, C) + hp = h_post.reshape(sb, n) + x_flat = x.reshape(sb, C) + + g_hr = torch.bmm(orig, go.transpose(1, 2)).view(s, b, n, n) + g_res = torch.bmm(hr, go).view(s, b, n, C) + g_x = torch.sum(go * hp.unsqueeze(-1), dim=1).view(s, b, C) + xb = x_flat if bias is None else x_flat + bias.view(1, C) + g_hp = torch.sum(go * xb.unsqueeze(1), dim=2).view(s, b, n) + g_bias = g_x.reshape(sb, C).sum(dim=0).to(dtype=bias.dtype) if bias is not None else None + return ( + g_hr.to(dtype=h_res.dtype), + g_res.to(dtype=original_residual.dtype), + g_hp.to(dtype=h_post.dtype), + g_x.to(dtype=x.dtype), + g_bias, + ) + + +@torch.compile +def _torch_proj_rms_compute_h( + x: Tensor, + weight: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + bias: Tensor, + n: int, + eps: float, + compute_h_eps: float = 1e-6, +) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + proj = torch.matmul(x, weight.t()) + r = x.norm(dim=-1, keepdim=True) / math.sqrt(x.shape[-1]) + alpha = torch.cat( + [alpha_pre.expand(n), alpha_post.expand(n), alpha_res.expand(weight.shape[0] - 2 * n)], + dim=-1, + ) + h = proj * alpha.unsqueeze(0) / (r + eps) + bias.unsqueeze(0) + h_pre = h[..., :n].sigmoid() + compute_h_eps + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res = h[..., 2 * n :] + return h_pre, h_post, h_res, r + + +if _CUTILE_AVAILABLE: + + class CutileSinkhornKnopp(torch.autograd.Function): + """cuTile Sinkhorn-Knopp projection fallback.""" + + @staticmethod + def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): + """Run cuTile Sinkhorn forward and save initial matrix for backward.""" + output, M_init = _cutile_sinkhorn_fwd(input_logits, num_iterations, eps) + ctx.save_for_backward(M_init) + ctx.num_iterations = num_iterations + ctx.eps = eps + return output + + @staticmethod + def backward(ctx, grad_output): + """Run cuTile Sinkhorn backward.""" + (M_init,) = ctx.saved_tensors + grad_input = _cutile_sinkhorn_bwd(grad_output, M_init, ctx.num_iterations, ctx.eps) + return grad_input, None, None + + class CutileHAggregate(torch.autograd.Function): + """cuTile n-stream weighted aggregation.""" + + @staticmethod + def forward(ctx, x: Tensor, h_pre: Tensor): + """Run cuTile h_aggregate forward.""" + output = _cutile_h_aggregate_fwd(x, h_pre) + ctx.save_for_backward(x, h_pre) + return output + + @staticmethod + def backward(ctx, grad_output): + """Run cuTile h_aggregate backward.""" + x, h_pre = ctx.saved_tensors + return _cutile_h_aggregate_bwd(grad_output, x, h_pre) + + class CutileProjRms(torch.autograd.Function): + """cuTile projection + RMS normalization.""" + + @staticmethod + def forward(ctx, x: Tensor, weight: Tensor, eps: float = 1e-6): + """Run cuTile projection plus RMS normalization forward.""" + proj, norm, r = _cutile_proj_rms_fwd(x, weight, eps) + ctx.save_for_backward(x, weight, norm) + ctx.eps = eps + return proj, r + + @staticmethod + def backward(ctx, grad_proj, grad_r): + """Run cuTile projection plus RMS normalization backward.""" + x, weight, norm = ctx.saved_tensors + grad_x, grad_weight = _cutile_proj_rms_bwd(grad_proj, grad_r, x, weight, norm, ctx.eps) + return grad_x, grad_weight, None + + class CutileProjRmsComputeH(torch.autograd.Function): + """cuTile projection + RMS norm + compute_h activations.""" + + @staticmethod + def forward( + ctx, + x: Tensor, + weight: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + bias: Tensor, + n: int, + eps: float = 1e-6, + compute_h_eps: float = 1e-6, + ): + """Run fused cuTile projection, RMS normalization, and compute_h forward.""" + h_pre, h_post, h_res, r, proj_reduced = _cutile_proj_rms_compute_h_fwd( + x, weight, bias, alpha_pre, alpha_post, alpha_res, n, eps, compute_h_eps + ) + ctx.save_for_backward( + x, + weight, + h_pre, + h_post, + h_res, + proj_reduced, + r, + alpha_pre, + alpha_post, + alpha_res, + bias, + ) + ctx.n = n + ctx.eps = eps + ctx.compute_h_eps = compute_h_eps + return h_pre, h_post, h_res, r + + @staticmethod + def backward(ctx, grad_h_pre, grad_h_post, grad_h_res, grad_r_ext): + """Run fused cuTile projection, RMS normalization, and compute_h backward.""" + ( + x, + weight, + h_pre, + h_post, + h_res, + proj, + r, + alpha_pre, + alpha_post, + alpha_res, + bias_param, + ) = ctx.saved_tensors + + grad_x, grad_weight, grad_ap, grad_apo, grad_ar, grad_bias = ( + _cutile_fused_compute_h_proj_rms_bwd( + x, + weight, + grad_h_pre, + grad_h_post, + grad_h_res, + h_pre, + h_post, + h_res, + proj, + r, + grad_r_ext, + alpha_pre, + alpha_post, + alpha_res, + bias_param, + ctx.n, + ctx.eps, + ctx.compute_h_eps, + ) + ) + + return (grad_x, grad_weight, grad_ap, grad_apo, grad_ar, grad_bias, None, None, None) + + +class FusedHAggregate(torch.autograd.Function): + """H_aggregate with Triton/cuTile/torch forward and cuTile/torch backward.""" + + @staticmethod + def forward(ctx, x: Tensor, h_pre: Tensor): + """Run h_aggregate forward using the best available backend.""" + triton_fwd = _get_triton_h_aggregate_fwd() + if triton_fwd is not None: + output = triton_fwd(x, h_pre) + elif is_cutile_available(): + output = _cutile_h_aggregate_fwd(x, h_pre) + else: + output = native_h_aggregate(x, h_pre) + ctx.save_for_backward(x, h_pre) + return output + + @staticmethod + def backward(ctx, grad_output): + """Run h_aggregate backward using the best available backend.""" + x, h_pre = ctx.saved_tensors + if is_cutile_available(): + return _cutile_h_aggregate_bwd(grad_output, x, h_pre) + return _torch_h_aggregate_bwd(grad_output, x, h_pre) + + +class FusedHPostBDA(torch.autograd.Function): + """H_post_bda with Triton/cuTile/torch forward and backward.""" + + @staticmethod + def forward( + ctx, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + x: Tensor, + bias: Optional[Tensor], + ): + """Run h_post_bda forward using the best available backend.""" + triton_fwd = _get_triton_h_post_bda_fwd() + if triton_fwd is not None: + output = triton_fwd(h_res, original_residual, h_post, x, bias) + elif is_cutile_available(): + output = _cutile_h_post_bda_fwd(h_res, original_residual, h_post, x, bias) + else: + output = native_h_post_bda(h_res, original_residual, h_post, x, bias) + if bias is not None: + ctx.save_for_backward(h_res, original_residual, h_post, x, bias) + ctx.has_bias = True + else: + ctx.save_for_backward(h_res, original_residual, h_post, x) + ctx.has_bias = False + return output + + @staticmethod + def backward(ctx, grad_output): + """Run h_post_bda backward using the best available backend.""" + if ctx.has_bias: + h_res, orig_res, h_post, x, bias = ctx.saved_tensors + else: + h_res, orig_res, h_post, x = ctx.saved_tensors + bias = None + + triton_bwd = _get_triton_h_post_bda_bwd() + if triton_bwd is not None: + return triton_bwd(grad_output, h_res, orig_res, h_post, x, bias) + if is_cutile_available(): + return _cutile_h_post_bda_bwd(grad_output, h_res, orig_res, h_post, x, bias) + return _torch_h_post_bda_bwd(grad_output, h_res, orig_res, h_post, x, bias) + + +def fused_sinkhorn(input_logits: Tensor, num_iterations: int, eps: float = 1e-6) -> Tensor: + """Project logits to a doubly stochastic matrix using Triton, cuTile, then torch.""" + _raise_mhc_backend_validation_error() + triton_sinkhorn = _get_triton_sinkhorn() + if triton_sinkhorn is not None: + return triton_sinkhorn(input_logits, num_iterations, eps) + if is_cutile_available(): + return CutileSinkhornKnopp.apply(input_logits, num_iterations, eps) + return native_sinkhorn(input_logits, num_iterations, eps) + + +def fused_h_aggregate(x: Tensor, h_pre: Tensor) -> Tensor: + """Weighted n-stream to 1-stream aggregation using Triton/cuTile/torch.""" + _raise_mhc_backend_validation_error() + if _TRITON_AVAILABLE or is_cutile_available(): + return FusedHAggregate.apply(x, h_pre) + return native_h_aggregate(x, h_pre) + + +def fused_h_post_bda( + h_res: Tensor, original_residual: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] +) -> Tensor: + """Fused H_res.T @ residual + H_post * (x + bias).""" + _raise_mhc_backend_validation_error() + if _TRITON_AVAILABLE or is_cutile_available(): + return FusedHPostBDA.apply(h_res, original_residual, h_post, x, bias) + return native_h_post_bda(h_res, original_residual, h_post, x, bias) + + +def fused_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6) -> Tuple[Tensor, Tensor]: + """Projection + RMS normalization using cuTile, then torch.""" + _raise_mhc_backend_validation_error() + if is_cutile_available(): + return CutileProjRms.apply(x, weight, eps) + return native_proj_rms(x, weight, eps) + + +def fused_proj_rms_compute_h( + x: Tensor, + weight: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + bias: Tensor, + n: int, + eps: float = 1e-6, + compute_h_eps: float = 1e-6, +) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + """Projection + RMS norm + compute_h split outputs using cuTile, then torch.""" + _raise_mhc_backend_validation_error() + if is_cutile_available(): + return CutileProjRmsComputeH.apply( + x, weight, alpha_pre, alpha_post, alpha_res, bias, n, eps, compute_h_eps + ) + return _torch_proj_rms_compute_h( + x, weight, alpha_pre, alpha_post, alpha_res, bias, n, eps, compute_h_eps + ) diff --git a/megatron/core/fusions/fused_mla_yarn_rope_apply.py b/megatron/core/fusions/fused_mla_yarn_rope_apply.py index 1fd5dcfae37..cf1c4a31fb0 100644 --- a/megatron/core/fusions/fused_mla_yarn_rope_apply.py +++ b/megatron/core/fusions/fused_mla_yarn_rope_apply.py @@ -29,17 +29,28 @@ @triton.jit def _get_thd_token_idx(cu_seqlens, pid_m, seq_num, cp_rank, cp_size): - token_idx = -1 - this_seq_len = 0 + # Cast ``pid_m`` and ``cu_seqlens`` loads to a single shared dtype so + # the loop-body reassignments don't surface as + # "initial value is int32 but redefined as int64" in newer Triton + # versions (which promote ``// Python_int`` to int64). + pid_m = pid_m.to(tl.int64) + token_idx = tl.full((), -1, dtype=tl.int64) + this_seq_len = tl.full((), 0, dtype=tl.int64) seq_idx = 0 - last_cum_seqlen = tl.load(cu_seqlens) // cp_size + last_cum_seqlen = tl.load(cu_seqlens).to(tl.int64) // cp_size while seq_idx < seq_num: - cur_cum_seqlen = tl.load(cu_seqlens + seq_idx + 1) // cp_size + cur_cum_seqlen = tl.load(cu_seqlens + seq_idx + 1).to(tl.int64) // cp_size if token_idx == -1 and cur_cum_seqlen > pid_m: token_idx = pid_m - last_cum_seqlen this_seq_len = cur_cum_seqlen - last_cum_seqlen last_cum_seqlen = cur_cum_seqlen seq_idx += 1 + # Padding tokens beyond cu_seqlens[-1] (from THD CUDA-graph padding) + # never match any sequence, leaving token_idx == -1. Clamp to 0 so + # the cos/sin table loads stay in-bounds; the wrong RoPE result is + # harmless because padding positions are excluded by loss_mask. + if token_idx == -1: + token_idx = tl.full((), 0, dtype=tl.int64) if cp_size > 1: if token_idx < this_seq_len // 2: token_idx = token_idx + cp_rank * this_seq_len // 2 @@ -65,29 +76,34 @@ def _get_thd_token_idx(cu_seqlens, pid_m, seq_num, cp_rank, cp_size): restore_value=["Q"], ) @triton.jit -def rotary_fwd_q_kernel( +def _mla_rope_fwd_inplace_kernel( Q, COS, SIN, - qk_head_dim, + nope_dim, emb_dim: tl.constexpr, head_num: tl.constexpr, batch_size, seq_num, cu_seqlens_q, + position_ids, stride_x_seq, stride_x_nheads, + stride_cos_seq, + stride_sin_seq, cp_rank, cp_size, + INVERSE: tl.constexpr, + REMOVE_INTERLEAVING: tl.constexpr, BLOCK_H: tl.constexpr, ): """ - Triton kernel of the forward pass for applying YARN RoPE to MLA's query. - This kernel inplace modifies the input tensor Q. + Forward pass: apply RoPE inplace to the trailing emb_dim elements. + Reads from interleaved layout, writes back to interleaved layout. Input: - Q: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] - or [total_seq_len, head_num, qk_head_dim + emb_dim] + Q: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] COS/SIN: [max_seq_len, emb_dim] batch_size: batch size for sbhd format, not used for thd format @@ -97,15 +113,24 @@ def rotary_fwd_q_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) - cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + cos_left = tl.load(COS + token_idx * stride_cos_seq + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * stride_sin_seq + tl.arange(0, emb_dim // 2)) + cos_right = tl.load( + COS + token_idx * stride_cos_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + sin_right = tl.load( + SIN + token_idx * stride_sin_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + if INVERSE: + sin_left = -sin_left + sin_right = -sin_right cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) @@ -113,7 +138,7 @@ def rotary_fwd_q_kernel( Q = Q + pid_m * stride_x_seq + pid_head * BLOCK_H * stride_x_nheads - x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + qk_head_dim + x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + nope_dim mask = x_off < head_num * stride_x_nheads # x1 = t[..., 0::2], x2 = t[..., 1::2] x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 @@ -124,10 +149,14 @@ def rotary_fwd_q_kernel( x_left = x_1 * cos_left - x_2 * sin_left x_right = x_2 * cos_right + x_1 * sin_right - x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] - x_right_off = x_left_off + emb_dim // 2 - tl.store(Q + x_left_off, x_left, mask=mask) - tl.store(Q + x_right_off, x_right, mask=mask) + if REMOVE_INTERLEAVING: + tl.store(Q + x_1_off, x_left, mask=mask) + tl.store(Q + x_2_off, x_right, mask=mask) + else: + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + tl.store(Q + x_left_off, x_left, mask=mask) + tl.store(Q + x_right_off, x_right, mask=mask) @triton.autotune( @@ -145,29 +174,34 @@ def rotary_fwd_q_kernel( restore_value=["DO"], ) @triton.jit -def rotary_bwd_q_kernel( +def _mla_rope_bwd_inplace_kernel( DO, COS, SIN, - qk_head_dim, + nope_dim, emb_dim: tl.constexpr, head_num: tl.constexpr, batch_size, seq_num, cu_seqlens_q, + position_ids, stride_x_seq, stride_x_nheads, + stride_cos_seq, + stride_sin_seq, cp_rank, cp_size, + INVERSE: tl.constexpr, + REMOVE_INTERLEAVING: tl.constexpr, BLOCK_H: tl.constexpr, ): """ - Triton kernel of the backward pass for applying YARN RoPE to MLA's query. - This kernel inplace modifies the input tensor DO. + Backward pass: inverse RoPE inplace on the trailing emb_dim elements. + Reads from interleaved layout, writes to interleaved layout. Input: - DO: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] - or [total_seq_len, head_num, qk_head_dim + emb_dim] + DO: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] COS/SIN: [max_seq_len, emb_dim] batch_size, seq_num, and cu_seqlens_q are the same as in the forward pass @@ -175,15 +209,24 @@ def rotary_bwd_q_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) - cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + cos_left = tl.load(COS + token_idx * stride_cos_seq + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * stride_sin_seq + tl.arange(0, emb_dim // 2)) + cos_right = tl.load( + COS + token_idx * stride_cos_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + sin_right = tl.load( + SIN + token_idx * stride_sin_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + if INVERSE: + sin_left = -sin_left + sin_right = -sin_right cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) @@ -191,25 +234,32 @@ def rotary_bwd_q_kernel( DO = DO + pid_m * stride_x_seq + pid_head * BLOCK_H * stride_x_nheads - x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + qk_head_dim + x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + nope_dim mask = x_off < head_num * stride_x_nheads - x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] - x_right_off = x_left_off + emb_dim // 2 - x_left = tl.load(DO + x_left_off, mask=mask) - x_right = tl.load(DO + x_right_off, mask=mask) + if REMOVE_INTERLEAVING: + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 + x_left = tl.load(DO + x_1_off, mask=mask) + x_right = tl.load(DO + x_2_off, mask=mask) + else: + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + x_left = tl.load(DO + x_left_off, mask=mask) + x_right = tl.load(DO + x_right_off, mask=mask) + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 x_1 = x_left * cos_left + x_right * sin_right x_2 = -x_left * sin_left + x_right * cos_right - x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 - x_2_off = x_1_off + 1 tl.store(DO + x_1_off, x_1, mask=mask) tl.store(DO + x_2_off, x_2, mask=mask) -class ApplyMLARotaryEmbQ(torch.autograd.Function): +class _FusedMLARoPEInplace(torch.autograd.Function): """ - Autograd function for applying YARN RoPE to MLA's query. + Autograd function for applying RoPE inplace to the trailing emb_dim + elements of a multi-head tensor (leaving the first nope_dim elements unchanged). """ @staticmethod @@ -218,22 +268,26 @@ def forward( q, cos, sin, - qk_head_dim, + nope_dim, emb_dim, cu_seqlens_q, cp_rank, cp_size, rotary_interleaved=False, + inverse=False, + remove_interleaving=False, + position_ids=None, ): """ - Forward function for ApplyMLARotaryEmbQ. + Forward function for _FusedMLARoPEInplace. Args: - q: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] - or [total_seq_len, head_num, qk_head_dim + emb_dim] + q: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] cos/sin: [max_seq_len, 1, 1, emb_dim] cu_seqlens_q: [seq_num + 1] accumulated sequence lengths for thd format rotary_interleaved: whether to apply RoPE interleaved, only supports False for now + inverse: if True, negate sin inside the kernel to apply the inverse rotation """ assert not rotary_interleaved max_seqlen = None @@ -241,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] @@ -248,33 +303,43 @@ 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.is_contiguous() - assert sin.is_contiguous() - assert headdim == qk_head_dim + emb_dim + assert cos.stride(-1) == 1 + assert sin.stride(-1) == 1 + assert headdim == nope_dim + emb_dim assert emb_dim % 4 == 0 grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_fwd_q_kernel[grid]( + _mla_rope_fwd_inplace_kernel[grid]( q, cos, sin, - qk_head_dim, + nope_dim, emb_dim, nheads, batch_size, seq_num, cu_seqlens_q, + position_ids, q.stride(0), q.stride(1), + cos.stride(0), + sin.stride(0), cp_rank, cp_size, + INVERSE=inverse, + REMOVE_INTERLEAVING=remove_interleaving, ) - ctx.save_for_backward(cos, sin) - ctx.qk_head_dim = qk_head_dim + 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 ctx.rotary_interleaved = rotary_interleaved + ctx.inverse = inverse + ctx.remove_interleaving = remove_interleaving ctx.cp_rank = cp_rank ctx.cp_size = cp_size if cu_seqlens_q is None: @@ -284,13 +349,17 @@ def forward( @staticmethod def backward(ctx, grad): """ - Backward function for ApplyMLARotaryEmbQ. + Backward function for _FusedMLARoPEInplace. Args: - grad: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] - or [total_seq_len, head_num, qk_head_dim + emb_dim] + 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 @@ -300,65 +369,126 @@ 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 grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_bwd_q_kernel[grid]( + _mla_rope_bwd_inplace_kernel[grid]( grad, cos, sin, - ctx.qk_head_dim, + ctx.nope_dim, ctx.emb_dim, nheads, batch_size, seq_num, ctx.cu_seqlens_q, + position_ids, grad.stride(0), grad.stride(1), + cos.stride(0), + sin.stride(0), ctx.cp_rank, ctx.cp_size, + INVERSE=ctx.inverse, + REMOVE_INTERLEAVING=ctx.remove_interleaving, ) 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 + return grad, None, None, None, None, None, None, None, None, None, None, None -def fused_apply_mla_rope_for_q( +def fused_mla_rope_inplace( t: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, - qk_head_dim: int, + nope_dim: int, emb_dim: int, cu_seqlens_q: Optional[torch.Tensor] = None, cp_rank: int = 0, cp_size: int = 1, rotary_interleaved: bool = False, -): + inverse: bool = False, + remove_interleaving: bool = False, + position_ids: Optional[torch.Tensor] = None, +) -> torch.Tensor: """ - Fused function for applying YARN RoPE to MLA's query. - This function inplace modifies the input tensor t. - Along the last dimension of t, the last emb_dim elements are applied with RoPE. - The first qk_head_dim elements are not modified. - It is an experimental feature and may change in future versions. + Fused RoPE applied inplace to the trailing emb_dim elements of a tensor, + leaving the first nope_dim elements unchanged. It supports both sbhd and thd input formats. + When ``inverse=True`` the rotation is reversed, which is useful for + undoing RoPE on the attention output. + For the notations below, seq_len is the length of the sequence per batch for sbhd format, total_seq_len is the total length of the sequences for thd format. max_seq_len is the maximum length of the sequences in the input tensor. Args: - t: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] - or [total_seq_len, head_num, qk_head_dim + emb_dim] + t: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] cos/sin: [max_seq_len, 1, 1, emb_dim] cu_seqlens_q: [seq_num + 1] accumulated sequence lengths for thd format 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 """ - return ApplyMLARotaryEmbQ.apply( - t, cos, sin, qk_head_dim, emb_dim, cu_seqlens_q, cp_rank, cp_size, rotary_interleaved + return _FusedMLARoPEInplace.apply( + t, + cos, + sin, + nope_dim, + emb_dim, + cu_seqlens_q, + cp_rank, + cp_size, + rotary_interleaved, + inverse, + remove_interleaving, + position_ids, + ) + + +def fused_mla_rope_out_of_place( + t: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + nope_dim: int, + emb_dim: int, + cu_seqlens_q: Optional[torch.Tensor] = None, + cp_rank: int = 0, + cp_size: int = 1, + rotary_interleaved: bool = False, + inverse: bool = False, + remove_interleaving: bool = False, + position_ids: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Apply the fused RoPE kernel without modifying the input tensor. + + Use this wrapper when an upstream autograd function may have retained its + output for backward. The underlying kernel remains in-place, so a private + copy is required to keep the retained tensor unchanged. + """ + return fused_mla_rope_inplace( + t.clone(), + cos, + sin, + nope_dim, + emb_dim, + cu_seqlens_q=cu_seqlens_q, + cp_rank=cp_rank, + cp_size=cp_size, + rotary_interleaved=rotary_interleaved, + inverse=inverse, + remove_interleaving=remove_interleaving, + position_ids=position_ids, ) @@ -376,7 +506,7 @@ def fused_apply_mla_rope_for_q( key=["emb_dim", "k_dim", "v_dim", "head_num"], ) @triton.jit -def rotary_fwd_kv_kernel( +def _mla_rope_fwd_kv_split_kernel( KV, K_POS_EMB, O_KEY, @@ -399,12 +529,12 @@ def rotary_fwd_kv_kernel( stride_v_nheads, cp_rank, cp_size, + REMOVE_INTERLEAVING: tl.constexpr, BLOCK_H: tl.constexpr, ): """ - Triton kernel of the forward pass for applying YARN RoPE to MLA's key and value. - It splits the input tensor KV into key and value, - and concatenates the processed RoPE to the key. + Forward pass: split KV into key and value, apply RoPE to k_pos_emb, + and concatenate the result onto key. Input: KV: [seq_len, batch_size, head_num, k_dim + v_dim] @@ -460,14 +590,24 @@ def rotary_fwd_kv_kernel( x_left = x_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) x_right = x_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - x_left_off = ( - tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads - + k_dim - + tl.arange(0, emb_dim // 2)[None, :] - ) - x_right_off = x_left_off + emb_dim // 2 - tl.store(K_ptr + x_left_off, x_left, mask=mask) - tl.store(K_ptr + x_right_off, x_right, mask=mask) + if REMOVE_INTERLEAVING: + x_1_off = ( + tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + + k_dim + + tl.arange(0, emb_dim // 2)[None, :] * 2 + ) + x_2_off = x_1_off + 1 + tl.store(K_ptr + x_1_off, x_left, mask=mask) + tl.store(K_ptr + x_2_off, x_right, mask=mask) + else: + x_left_off = ( + tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + + k_dim + + tl.arange(0, emb_dim // 2)[None, :] + ) + x_right_off = x_left_off + emb_dim // 2 + tl.store(K_ptr + x_left_off, x_left, mask=mask) + tl.store(K_ptr + x_right_off, x_right, mask=mask) @triton.autotune( @@ -484,7 +624,7 @@ def rotary_fwd_kv_kernel( key=["emb_dim", "k_dim", "v_dim", "head_num"], ) @triton.jit -def rotary_bwd_kv_kernel( +def _mla_rope_bwd_kv_split_kernel( dK, dV, dKV, @@ -507,10 +647,11 @@ def rotary_bwd_kv_kernel( stride_demb_seq, cp_rank, cp_size, + REMOVE_INTERLEAVING: tl.constexpr, BLOCK_H: tl.constexpr, ): """ - Triton kernel of the backward pass for applying YARN RoPE to MLA's key and value. + Backward pass for the KV-split RoPE. Input: dK: [seq_len, batch_size, head_num, emb_dim + k_dim] @@ -555,10 +696,16 @@ def rotary_bwd_kv_kernel( dK_ptr = dK + pid_m * stride_dk_seq + i * BLOCK_H * stride_dk_nheads x_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + k_dim mask = x_off < head_num * stride_dk_nheads - x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] - x_right_off = x_left_off + emb_dim // 2 - x_left = tl.load(dK_ptr + x_left_off, mask=mask) - x_right = tl.load(dK_ptr + x_right_off, mask=mask) + if REMOVE_INTERLEAVING: + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 + x_left = tl.load(dK_ptr + x_1_off, mask=mask) + x_right = tl.load(dK_ptr + x_2_off, mask=mask) + else: + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + x_left = tl.load(dK_ptr + x_left_off, mask=mask) + x_right = tl.load(dK_ptr + x_right_off, mask=mask) x_left_accum += x_left x_right_accum += x_right x_left_accum = tl.sum(x_left_accum, axis=0) @@ -578,9 +725,10 @@ def rotary_bwd_kv_kernel( tl.store(dEMB_ptr + tl.arange(0, emb_dim // 2) * 2 + 1, x_2) -class ApplyMLARotaryEmbKV(torch.autograd.Function): +class _FusedMLARoPEKVSplit(torch.autograd.Function): """ - Autograd function for applying YARN RoPE to MLA's key and value. + Autograd function for applying RoPE to MLA's key and value. + Splits KV, applies RoPE to k_pos_emb, concatenates onto key. """ @staticmethod @@ -597,9 +745,10 @@ def forward( cp_rank, cp_size, rotary_interleaved=False, + remove_interleaving=False, ): """ - Forward function for ApplyMLARotaryEmbKV. + Forward function for _FusedMLARoPEKVSplit. Args: kv: [seq_len, batch_size, head_num, k_dim + v_dim] @@ -634,7 +783,7 @@ def forward( o_value = kv.new_empty(total_seqlen, nheads, v_dim) grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_fwd_kv_kernel[grid]( + _mla_rope_fwd_kv_split_kernel[grid]( kv, k_pos_emb, o_key, @@ -657,8 +806,10 @@ def forward( o_value.stride(1), cp_rank, cp_size, + REMOVE_INTERLEAVING=remove_interleaving, ) ctx.save_for_backward(cos, sin) + ctx.remove_interleaving = remove_interleaving ctx.rotary_interleaved = rotary_interleaved ctx.emb_dim = emb_dim ctx.k_dim = k_dim @@ -674,7 +825,7 @@ def forward( @staticmethod def backward(ctx, dk, dv): """ - Backward function for ApplyMLARotaryEmbKV. + Backward function for _FusedMLARoPEKVSplit. Args: dk: [seq_len, batch_size, head_num, emb_dim + k_dim] @@ -702,7 +853,7 @@ def backward(ctx, dk, dv): d_emb = dk.new_empty(total_seqlen, 1, ctx.emb_dim) grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_bwd_kv_kernel[grid]( + _mla_rope_bwd_kv_split_kernel[grid]( dk, dv, d_kv, @@ -725,14 +876,15 @@ def backward(ctx, dk, dv): d_emb.stride(0), ctx.cp_rank, ctx.cp_size, + REMOVE_INTERLEAVING=ctx.remove_interleaving, ) if ctx.cu_seqlens_kv is None: d_kv = d_kv.view(max_seqlen, batch_size, nheads, ctx.k_dim + ctx.v_dim) d_emb = d_emb.view(max_seqlen, batch_size, 1, ctx.emb_dim) - return d_kv, d_emb, None, None, None, None, None, None, None, None, None + return d_kv, d_emb, None, None, None, None, None, None, None, None, None, None -def fused_apply_mla_rope_for_kv( +def fused_mla_rope_kv_split( kv: torch.Tensor, k_pos_emb: torch.Tensor, cos: torch.Tensor, @@ -744,9 +896,10 @@ def fused_apply_mla_rope_for_kv( cp_rank: int = 0, cp_size: int = 1, rotary_interleaved: bool = False, -): + remove_interleaving: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: """ - Fused function for applying YARN RoPE to MLA's key and value. + Fused function for applying RoPE to MLA's key and value. It splits the input tensor kv into key and value, and concatenates the processed RoPE to the key. @@ -761,13 +914,14 @@ def fused_apply_mla_rope_for_kv( cos/sin: [max_seq_len, 1, 1, emb_dim] cu_seqlens_kv: [seq_num + 1] accumulated sequence lengths for thd format rotary_interleaved: whether to apply RoPE interleaved, only supports False for now + remove_interleaving: if True, output RoPE dims in non-interleaved layout Returns: key: [seq_len, batch_size, head_num, emb_dim + k_dim] or [total_seq_len, head_num, emb_dim + k_dim] value: [seq_len, batch_size, head_num, v_dim] or [total_seq_len, head_num, v_dim] """ - return ApplyMLARotaryEmbKV.apply( + return _FusedMLARoPEKVSplit.apply( kv, k_pos_emb, cos, @@ -779,4 +933,64 @@ def fused_apply_mla_rope_for_kv( cp_rank, cp_size, rotary_interleaved, + remove_interleaving, + ) + + +def fused_apply_mla_rope_for_q( + t: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + qk_head_dim: int, + emb_dim: int, + cu_seqlens_q: Optional[torch.Tensor] = None, + cp_rank: int = 0, + cp_size: int = 1, + rotary_interleaved: bool = False, +) -> torch.Tensor: + """Backward-compatible in-place MLA query RoPE API. + + New callers should choose :func:`fused_mla_rope_inplace` or + :func:`fused_mla_rope_out_of_place` explicitly. This legacy name keeps + its original mutation behavior and does not add a clone to the hot path. + """ + return fused_mla_rope_inplace( + t, + cos, + sin, + qk_head_dim, + emb_dim, + cu_seqlens_q=cu_seqlens_q, + cp_rank=cp_rank, + cp_size=cp_size, + rotary_interleaved=rotary_interleaved, + ) + + +def fused_apply_mla_rope_for_kv( + kv: torch.Tensor, + k_pos_emb: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + emb_dim: int, + k_dim: int, + v_dim: int, + cu_seqlens_kv: Optional[torch.Tensor] = None, + cp_rank: int = 0, + cp_size: int = 1, + rotary_interleaved: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Backward-compatible name for the MLA key/value split RoPE API.""" + return fused_mla_rope_kv_split( + kv, + k_pos_emb, + cos, + sin, + emb_dim, + k_dim, + v_dim, + cu_seqlens_kv=cu_seqlens_kv, + cp_rank=cp_rank, + cp_size=cp_size, + rotary_interleaved=rotary_interleaved, ) diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py index 88bb070e105..3d6b053be60 100644 --- a/megatron/core/model_parallel_config.py +++ b/megatron/core/model_parallel_config.py @@ -2,11 +2,23 @@ import warnings from dataclasses import dataclass, field -from typing import Callable, ContextManager, Literal, Optional +from typing import Callable, ContextManager, Literal, Optional, Union import torch +def _parse_pad_packed_seq_alignment(value): + """Parse THD packed-sequence padding alignment.""" + if value == "max": + return value + try: + return int(value) + except (TypeError, ValueError) as exc: + raise ValueError( + "pad_packed_seq_alignment must be 'max' or a positive integer alignment." + ) from exc + + @dataclass class ModelParallelConfig: """Base configuration for Megatron Core @@ -59,7 +71,7 @@ class ModelParallelConfig: can handle without overflowing the memory. Typically, a good starting point is to set this to maximum sequence length / context parallel size. This is used to calculate the number and length of sub-samples assigned to - each rank when using hybrid_context_parallel. + each rank when hybrid_context_parallel or sequence_packing_scheduler is enabled. """ hybrid_context_parallel: bool = False @@ -69,6 +81,27 @@ class ModelParallelConfig: Please set max_seqlen_per_dp_cp_rank when using hybrid_context_parallel. """ + sequence_packing_scheduler: Optional[Literal['dp_balanced']] = None + """ + Scheduler for packing variable-length THD batches. + dp_balanced: DP-balanced scheduler for sequence packing. + """ + + pad_packed_seq_alignment: Optional[Union[int, Literal["max"]]] = field( + default=None, + metadata={ + "argparse_meta": { + "arg_names": ["--pad-packed-seq-alignment"], + "type": _parse_pad_packed_seq_alignment, + } + }, + ) + """Pad packed THD tensors to ``max_seqlen_per_dp_cp_rank`` (``"max"``) + or to a positive integer alignment after packing.""" + + pad_packed_seq_by_appending_dummy_seq: bool = True + """Represent the post-pack padding tail as a dummy THD sequence.""" + expert_model_parallel_size: int = 1 """Distributes Moe Experts across sub data parallel dimension.""" @@ -423,6 +456,27 @@ def __post_init__(self): See https://docs.python.org/3/library/dataclasses.html#post-init-processing for more details. """ + if self.pad_packed_seq_alignment is not None: + self.pad_packed_seq_alignment = _parse_pad_packed_seq_alignment( + self.pad_packed_seq_alignment + ) + if self.max_seqlen_per_dp_cp_rank is None: + raise ValueError( + "max_seqlen_per_dp_cp_rank must be set when " + "pad_packed_seq_alignment is enabled." + ) + if self.pad_packed_seq_alignment != "max": + if self.pad_packed_seq_alignment <= 0: + raise ValueError( + "pad_packed_seq_alignment must be 'max' or a positive integer alignment." + ) + if self.pad_packed_seq_alignment > self.max_seqlen_per_dp_cp_rank: + raise ValueError( + "pad_packed_seq_alignment must not exceed " + f"max_seqlen_per_dp_cp_rank ({self.max_seqlen_per_dp_cp_rank}), " + f"got {self.pad_packed_seq_alignment}." + ) + if self.sequence_parallel: if self.tensor_model_parallel_size <= 1: raise ValueError("Cannot use sequence parallelism without tensor parallelism") diff --git a/megatron/core/models/common/embeddings/rope_utils.py b/megatron/core/models/common/embeddings/rope_utils.py index 9fab25a3fae..0d6b5ea5fb5 100644 --- a/megatron/core/models/common/embeddings/rope_utils.py +++ b/megatron/core/models/common/embeddings/rope_utils.py @@ -95,6 +95,8 @@ def _apply_rotary_pos_emb_bshd( rotary_interleaved: bool = False, mla_rotary_interleaved: bool = False, mscale: float = 1.0, + inverse: bool = False, + mla_output_remove_interleaving: bool = False, multi_latent_attention: Optional[bool] = None, ) -> Tensor: """Apply rotary positional embedding to input tensor T. @@ -118,6 +120,13 @@ def _apply_rotary_pos_emb_bshd( ) mla_rotary_interleaved = multi_latent_attention + # Some callers may pass freqs with an extra singleton axis, e.g. + # t: [s, b, d] and freqs: [s, 1, 1, d]. In that case, broadcasting would + # accidentally expand to [s, s, b, d]. Squeeze the extra singleton axis to + # keep freqs rank aligned with t. + if freqs.dim() == t.dim() + 1 and freqs.size(-2) == 1: + freqs = freqs.squeeze(-2) + rot_dim = freqs.shape[-1] # ideally t_pass is empty so rotary pos embedding is applied to all tensor t @@ -132,8 +141,18 @@ def _apply_rotary_pos_emb_bshd( # second part is sine component, need to change signs with _rotate_half method cos_ = (torch.cos(freqs) * mscale).to(t.dtype) sin_ = (torch.sin(freqs) * mscale).to(t.dtype) + if inverse: + sin_ = -sin_ t = (t * cos_) + (_rotate_half(t, rotary_interleaved) * sin_) + + # Fallback to original permutation + # DSv4 applies rope on V and O, so we need to uninterleave the tensor. + # The existing MLA code is safe because the dot product is permutation-invariant. + if mla_rotary_interleaved and mla_output_remove_interleaving: + x1, x2 = torch.chunk(t, 2, dim=-1) + t = torch.stack((x1, x2), dim=-1).flatten(start_dim=-2) + return torch.cat((t, t_pass), dim=-1) @@ -193,20 +212,27 @@ def _apply_rotary_pos_emb_thd( rotary_interleaved: bool = False, mla_rotary_interleaved: bool = False, mscale: float = 1.0, + inverse: bool = False, + mla_output_remove_interleaving: bool = False, cp_group: torch.distributed.ProcessGroup = None, multi_latent_attention: Optional[bool] = None, + max_seqlen: Optional[int] = None, ) -> Tensor: - """A baseline implementation of applying RoPE for `thd` format. + """Apply RoPE for `thd` format using pure CUDA ops (CUDA Graph compatible). + + Replaces the original Python-loop + .tolist() implementation with vectorized + CUDA operations. No GPU->CPU syncs, compatible with CUDA Graph capture. Args: - t (Tensor): Input tensor T is of shape [t, h, d] - cu_seqlens(Tensor): Cumulative sum of sequence lengths in a batch for `t`, - with shape [b + 1] and dtype torch.int32. - freqs (Tensor): Rotary Positional embedding tensor freq is of shape [max_s, 1, 1, d] - cp_group (torch.distributed.ProcessGroup): The context parallel group + t (Tensor): Input tensor of shape [total_tokens, h, d] + cu_seqlens (Tensor): Cumulative sequence lengths, shape [num_seqs + 1], int32. + freqs (Tensor): RoPE frequencies, shape [max_s, 1, 1, d] or [total_tokens, 1, 1, d] + cp_group: Context parallel group + max_seqlen: Global max sequence length for this packed batch when known. Supplying it + avoids the compatibility-path host sync used by legacy callers. Returns: - Tensor: Shape [t, h, d]. The input tensor after applying RoPE. + Tensor: Shape [total_tokens, h, d]. Input with RoPE applied. """ if multi_latent_attention is not None: warnings.warn( @@ -219,53 +245,71 @@ def _apply_rotary_pos_emb_thd( raise ValueError("cp_group must be provided for THD format RoPE") cp_size = cp_group.size() cp_rank = cp_group.rank() - seqlens = ((cu_seqlens[1:] - cu_seqlens[:-1]) // cp_size).tolist() - sequence_splits = torch.split(t, seqlens) - total_seqlen = int(cu_seqlens[-1].item()) - has_packed_freqs = freqs.dim() >= 1 and freqs.size(0) == total_seqlen - - # Handle two different frequency tensor formats: - # 1. If freqs.size(0) == cu_seqlens[-1]: freqs contains positions for the whole packed - # batch. Each sequence must therefore use its cu_seqlens offset when selecting the local CP - # front/back slices. For example, with cu_seqlens=[0, 4, 8], cp_size=2, rank 0 should use - # positions [0, 3, 4, 7], not [0, 3, 0, 3]. - # 2. Otherwise: freqs contains only max sequence length positions. Each packed sequence should - # reuse positions starting from 0, preserving the legacy THD behavior. - if has_packed_freqs: - # CASE 1: Exact mapping with offsets - local_freqs = [] - for i, x in enumerate(sequence_splits): - # cu_seqlens[i] is the starting offset of this sequence in the original batch - seq_start_offset = cu_seqlens[i].item() - local_freqs.append( - _get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs, seq_start_offset) - ) - freqs = torch.cat(local_freqs, dim=0) - return _apply_rotary_pos_emb_bshd( - t.unsqueeze(1), - freqs, - rotary_interleaved=rotary_interleaved, - mla_rotary_interleaved=mla_rotary_interleaved, - mscale=mscale, - ).squeeze(1) - - # CASE 2: Traditional mapping without offsets. Apply RoPE one sequence at a time so the second - # and later packed sequences do not look like continuations of the first sequence. - output = torch.empty_like(t) - output_offset = 0 - for x in sequence_splits: - freq_slice = _get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs) - output_slice = _apply_rotary_pos_emb_bshd( - x.unsqueeze(1), - freq_slice, - rotary_interleaved=rotary_interleaved, - mla_rotary_interleaved=mla_rotary_interleaved, - mscale=mscale, - ).squeeze(1) - output.narrow(0, output_offset, x.size(0)).copy_(output_slice) - output_offset += x.size(0) + total_tokens = t.shape[0] + device = t.device + + token_pos = torch.arange(total_tokens, device=device, dtype=torch.int64) + + # `cu_seqlens` describes the global packed sequence. With CP, `t` is already + # CP-partitioned, so build a local cumulative-length view before assigning + # local tokens to packed sequences. + cu_seqlens_i64 = cu_seqlens.to(torch.int64) + global_seq_lens = cu_seqlens_i64[1:] - cu_seqlens_i64[:-1] + local_seq_lens = global_seq_lens // cp_size if cp_size > 1 else global_seq_lens + local_cu_seqlens = torch.zeros_like(cu_seqlens_i64) + local_cu_seqlens[1:] = torch.cumsum(local_seq_lens, dim=0) + + # `searchsorted(..., right=True) - 1` returns the local sequence index. The + # clamp guards padded tokens that sit beyond the final real local token; they + # get a harmless frequency and are later masked out. + seq_idx = torch.searchsorted(local_cu_seqlens, token_pos, right=True) - 1 + seq_idx = seq_idx.clamp(min=0, max=cu_seqlens.shape[0] - 2) + + local_seq_start = local_cu_seqlens[seq_idx] + local_pos = token_pos - local_seq_start + local_seq_len = local_seq_lens[seq_idx] + global_seq_start = cu_seqlens_i64[seq_idx] - return output + if cp_size > 1: + cp_seg = local_seq_len // 2 + full_seqlen = local_seq_len * cp_size + is_first_half = local_pos < cp_seg + freq_pos = torch.where( + is_first_half, + cp_rank * cp_seg + local_pos, + full_seqlen - (cp_rank + 1) * cp_seg + (local_pos - cp_seg), + ) + else: + freq_pos = local_pos.to(torch.int64) + + if max_seqlen is None: + # Backward compatibility for callers that predate ``max_seqlen``. This retains + # the old packed-frequency semantics at the cost of a GPU-to-CPU sync. Updated + # training paths pass ``max_seqlen`` and stay CUDA-graph safe. + exact_packed_freqs = freqs.dim() >= 1 and freqs.size(0) == int(cu_seqlens[-1].item()) + else: + exact_packed_freqs = freqs.dim() >= 1 and freqs.size(0) > max_seqlen + if exact_packed_freqs: + # `freqs` covers all positions across all sequences (used for non-1D + # RoPE / VLMs); shift by the per-sequence start offset so each token + # samples its absolute position. When `freqs` only spans one max-len + # sequence, no shift is needed. + freq_pos = freq_pos + global_seq_start + + # Padded positions can sit outside the frequency table. Clamp them into + # range; downstream padding masks exclude those positions from the result. + freq_pos = freq_pos.clamp(min=0, max=freqs.shape[0] - 1) + freqs_packed = freqs[freq_pos] + + return _apply_rotary_pos_emb_bshd( + t.unsqueeze(1), + freqs_packed, + rotary_interleaved=rotary_interleaved, + mla_rotary_interleaved=mla_rotary_interleaved, + mscale=mscale, + inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, + ).squeeze(1) def apply_rotary_pos_emb( @@ -276,6 +320,9 @@ def apply_rotary_pos_emb( mscale: float = 1.0, cp_group: torch.distributed.ProcessGroup = None, mla_rotary_interleaved: bool = False, + inverse: bool = False, + mla_output_remove_interleaving: bool = False, + max_seqlen: Optional[int] = None, ): """ Reroute to the appropriate apply_rotary_pos_emb function depending on @@ -312,6 +359,12 @@ def apply_rotary_pos_emb( "Using unfused implementation." ) use_unfused = True + if inverse: + warnings.warn( + "inverse RoPE is not supported by TE's fused RoPE. " + "Using unfused implementation." + ) + use_unfused = True if not use_unfused: assert fused_apply_rotary_pos_emb is not None, "apply_rope_fusion is not available." return fused_apply_rotary_pos_emb(t, freqs, interleaved=config.rotary_interleaved) @@ -333,6 +386,8 @@ def apply_rotary_pos_emb( rotary_interleaved=config.rotary_interleaved, mla_rotary_interleaved=mla_rotary_interleaved, mscale=mscale, + inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, ) else: return _apply_rotary_pos_emb_thd( @@ -343,6 +398,9 @@ def apply_rotary_pos_emb( mla_rotary_interleaved=mla_rotary_interleaved, mscale=mscale, cp_group=cp_group, + inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, + max_seqlen=max_seqlen, ) diff --git a/megatron/core/models/common/embeddings/rotary_pos_embedding.py b/megatron/core/models/common/embeddings/rotary_pos_embedding.py index 0e560f939f2..e591e4ff90d 100644 --- a/megatron/core/models/common/embeddings/rotary_pos_embedding.py +++ b/megatron/core/models/common/embeddings/rotary_pos_embedding.py @@ -205,6 +205,47 @@ def forward( return emb + def _set_cos_sin_cache(self, seq_len, offset, dtype, packed_seq=False, cp_group=None): + """Materialize cached cos/sin tensors for ``[seq_len, ..., dim]``.""" + self.max_seq_len_cached = seq_len + self.offset_cached = offset + self.dtype_cached = dtype + self.packed_seq_cached = packed_seq + + emb = self.forward(seq_len, offset, packed_seq=packed_seq, cp_group=cp_group) + self.register_buffer("cos_cached", emb.cos().to(dtype).contiguous(), persistent=False) + self.register_buffer("sin_cached", emb.sin().to(dtype).contiguous(), persistent=False) + + def get_cached_cos_sin( + self, + seq_len, + offset=0, + dtype=torch.get_default_dtype(), + packed_seq=False, + cp_group=None, + mscale=None, + ): + """Get cached cos and sin values. + + The cache is rebuilt on first use or whenever ``seq_len`` grows + beyond the cached length, or any of ``offset`` / ``dtype`` / + ``packed_seq`` changes from the previous call. + ``YarnRotaryEmbedding`` overrides this to also bake its + concentration factor into the cached cos/sin (controlled by + ``mscale``); for the base class without a concentration + factor the argument is accepted-and-ignored for API uniformity. + """ + del mscale # base class has no concentration factor + if ( + not hasattr(self, "max_seq_len_cached") + or seq_len > self.max_seq_len_cached + or offset != self.offset_cached + or dtype != self.dtype_cached + or packed_seq != self.packed_seq_cached + ): + self._set_cos_sin_cache(seq_len, offset, dtype, packed_seq, cp_group) + return (self.cos_cached[:seq_len, ...], self.sin_cached[:seq_len, ...]) + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): state_dict.pop(f'{prefix}inv_freq', None) return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) diff --git a/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py b/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py index 166ef9b41e7..cb8a03d0b2b 100644 --- a/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py +++ b/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py @@ -186,13 +186,18 @@ def forward( emb = get_pos_emb_on_this_cp_rank(emb, 0, cp_group) return emb, _mscale - def _set_cos_sin_cache(self, seq_len, offset, dtype, packed_seq=False): + def _set_cos_sin_cache( + self, seq_len, offset, dtype, packed_seq=False, cp_group=None, mscale=None + ): self.max_seq_len_cached = seq_len self.offset_cached = offset self.dtype_cached = dtype self.packed_seq_cached = packed_seq + self.mscale_cached = mscale - emb, _mscale = self.forward(seq_len, offset, packed_seq) + emb, _mscale = self.forward(seq_len, offset, packed_seq=packed_seq, cp_group=cp_group) + if mscale is not None: + _mscale = mscale self.register_buffer( "cos_cached", (emb.cos() * _mscale).to(dtype).contiguous(), persistent=False ) @@ -201,16 +206,34 @@ def _set_cos_sin_cache(self, seq_len, offset, dtype, packed_seq=False): ) def get_cached_cos_sin( - self, seq_len, offset=0, dtype=torch.get_default_dtype(), packed_seq=False + self, + seq_len, + offset=0, + dtype=torch.get_default_dtype(), + packed_seq=False, + cp_group=None, + mscale=None, ): - """Get cached cos and sin values.""" + """Get cached cos and sin values. + + Args: + mscale: when ``None`` (default), the cached cos/sin are + multiplied by yarn's internal concentration factor (the + normal long-context behaviour). When a float is supplied, + that value is used in place of the internal factor — e.g. + the DSv4 hybrid model passes ``mscale=1.0`` to enforce + its "pure rotation" contract and keep the fused / + unfused rope paths bit-equivalent. + """ if ( - seq_len > self.max_seq_len_cached + not hasattr(self, "max_seq_len_cached") + or seq_len > self.max_seq_len_cached or offset != self.offset_cached or dtype != self.dtype_cached or packed_seq != self.packed_seq_cached + or mscale != getattr(self, "mscale_cached", None) ): - self._set_cos_sin_cache(seq_len, offset, dtype, packed_seq) + self._set_cos_sin_cache(seq_len, offset, dtype, packed_seq, cp_group, mscale) return (self.cos_cached[:seq_len, ...], self.sin_cached[:seq_len, ...]) diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py index 22322e1b346..7471d5fdecd 100644 --- a/megatron/core/models/hybrid/hybrid_block.py +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -7,7 +7,7 @@ from contextlib import nullcontext from dataclasses import dataclass -from typing import Optional, Tuple, Union +from typing import List, Optional, Tuple, Union import torch from torch import Tensor, nn @@ -21,16 +21,31 @@ from megatron.core.inference.contexts import BaseInferenceContext from megatron.core.inference.utils import InferenceMode from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols as LayerSymbols -from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.packed_seq_params import PackedSeqParams, has_packed_seq_params_cuda_graph_kwargs from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.recompute import checkpointed_forward +from megatron.core.tensor_parallel.random import CheckpointManager from megatron.core.transformer import TransformerConfig from megatron.core.transformer.cuda_graphs import annotate_first_last_layer +from megatron.core.transformer.enums import CudaGraphModule +from megatron.core.transformer.hyper_connection import ( + HyperConnectionModule, + learned_output_contract, +) from megatron.core.transformer.identity_op import IdentityOp -from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.module import ( + GraphableMegatronModule, + MegatronModule, + convert_module_to_dtype_except_fp32_marked, + mark_keep_in_fp32, +) from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.transformer_layer import TransformerLayer -from megatron.core.transformer.utils import sharded_state_dict_default +from megatron.core.transformer.utils import ( + ensure_metadata_has_dp_cp_group, + make_sharded_tensors_for_checkpoint, + sharded_state_dict_default, +) from megatron.core.utils import WrappedTensor, deprecate_inference_params, make_viewless_tensor @@ -44,11 +59,465 @@ class HybridStackSubmodules: gdn_layer: Union[ModuleSpec, type] = IdentityOp attention_layer: Union[ModuleSpec, type] = IdentityOp dsa_layer: Union[ModuleSpec, type] = IdentityOp + csa_layer: Union[ModuleSpec, type] = IdentityOp + hca_layer: Union[ModuleSpec, type] = IdentityOp + window_layer: Union[ModuleSpec, type] = IdentityOp mlp_layer: Union[ModuleSpec, type] = IdentityOp moe_layer: Union[ModuleSpec, type] = IdentityOp mtp_block_spec: Optional[ModuleSpec] = None +class HyperConnectionHybridLayer(GraphableMegatronModule): + """Layer-boundary mHC wrapper for HybridStack layers. + + Hybrid layers already own their local residual paths. Each wrapped layer is + treated as one function by aggregating n streams to its input, running the + existing layer, and feeding only the layer delta back through mHC expansion. + + This wrapper nests the inner layer under inner_layer. Checkpoints cannot + switch between mHC-enabled and ordinary HybridStacks without key migration. + """ + + def __init__(self, config: TransformerConfig, layer: MegatronModule) -> None: + super().__init__(config=config) + self.inner_layer = layer + self.layer_number = layer.layer_number + self.offload_module_in_cuda_graph = getattr( + layer, "offload_module_in_cuda_graph", False + ) + self.hyper_connection = HyperConnectionModule(config=config, layer_number=self.layer_number) + if config.params_dtype is not None: + convert_module_to_dtype_except_fp32_marked( + self.hyper_connection, config.params_dtype + ) + if hasattr(layer, 'tp_group'): + self.tp_group = layer.tp_group + + def create_mcore_cudagraph_manager(self, config): + """Leave local CUDA graph routing on the already-constructed inner layer. + + This wrapper adds a TE per-layer graph boundary. Installing a second + local manager here would capture the wrapper regardless of the inner + layer's local graph scope and could nest the inner manager. + """ + return None + + def get_layer_static_inputs(self, seq_length, micro_batch_size): + """Return the inner layer's static inputs with n-stream hidden width.""" + if hasattr(self.inner_layer, "get_layer_static_inputs"): + static_inputs = self.inner_layer.get_layer_static_inputs( + seq_length, micro_batch_size + ) + else: + static_inputs = super().get_layer_static_inputs( + seq_length, micro_batch_size + ) + hidden_states = static_inputs["hidden_states"] + static_inputs["hidden_states"] = torch.ones( + ( + hidden_states.shape[0], + hidden_states.shape[1], + self.config.num_residual_streams * self.config.hidden_size, + ), + dtype=hidden_states.dtype, + requires_grad=hidden_states.requires_grad, + device=hidden_states.device, + ) + return static_inputs + + def _set_te_cuda_graph_packed_seq_params_static_metadata( + self, static_metadata, tensor_kwarg_names=None + ): + """Delegate the #5672 packed-sequence graph contract to the inner layer.""" + return self.inner_layer._set_te_cuda_graph_packed_seq_params_static_metadata( + static_metadata, tensor_kwarg_names + ) + + def _get_te_cuda_graph_packed_seq_params_static_metadata(self): + return self.inner_layer._get_te_cuda_graph_packed_seq_params_static_metadata() + + def _validate_te_cuda_graph_packed_seq_params_static_metadata( + self, static_metadata + ): + return self.inner_layer._validate_te_cuda_graph_packed_seq_params_static_metadata( + static_metadata + ) + + def _get_te_cuda_graph_packed_seq_params_tensor_kwarg_names(self): + return self.inner_layer._get_te_cuda_graph_packed_seq_params_tensor_kwarg_names() + + def _validate_te_cuda_graph_packed_seq_params_tensor_kwargs(self, tensor_kwargs): + return self.inner_layer._validate_te_cuda_graph_packed_seq_params_tensor_kwargs( + tensor_kwargs + ) + + def _rebuild_te_cuda_graph_packed_seq_params(self, kwargs): + if hasattr(self.inner_layer, "_rebuild_te_cuda_graph_packed_seq_params"): + return self.inner_layer._rebuild_te_cuda_graph_packed_seq_params(kwargs) + assert not has_packed_seq_params_cuda_graph_kwargs(kwargs), ( + "Wrapped non-Transformer layers do not support flattened " + "PackedSeqParams CUDA graph inputs." + ) + return None + + def _flatten_te_cuda_graph_packed_seq_params(self, kwargs): + if hasattr(self.inner_layer, "_flatten_te_cuda_graph_packed_seq_params"): + if ( + self.inner_layer._get_te_cuda_graph_packed_seq_params_static_metadata() + is None + ): + assert not has_packed_seq_params_cuda_graph_kwargs(kwargs), ( + "Wrapped Transformer layers captured without packed-sequence " + "metadata cannot receive flattened PackedSeqParams CUDA graph inputs." + ) + kwargs.pop("packed_seq_params", None) + return None + return self.inner_layer._flatten_te_cuda_graph_packed_seq_params(kwargs) + assert kwargs.get("packed_seq_params") is None, ( + "Wrapped non-Transformer layers do not support PackedSeqParams " + "CUDA graph replay." + ) + kwargs.pop("packed_seq_params", None) + return None + + def __call__(self, *args, **kwargs): + """Keep the non-Tensor mHC recompute manager outside TE graph inputs.""" + self._mhc_recompute_manager = kwargs.pop("mhc_recompute_manager", None) + try: + return super().__call__(*args, **kwargs) + finally: + self._mhc_recompute_manager = None + + def _inner_is_moe(self) -> bool: + from megatron.core.transformer.moe.moe_layer import MoELayer + + return isinstance(self.inner_layer, TransformerLayer) and isinstance( + getattr(self.inner_layer, 'mlp', None), MoELayer + ) + + def _inner_is_partial_moe_capture(self) -> bool: + return ( + self._inner_is_moe() + and bool(self.config.cuda_graph_modules) + and CudaGraphModule.moe_router in self.config.cuda_graph_modules + ) + + def _te_cuda_graph_capture(self, *args, **kwargs): + """Capture the whole wrapper, or only its graph-safe partial-MoE prefix.""" + sample_kwarg_names = frozenset(kwargs) + captured_sample_kwarg_names = getattr( + self, "_te_cuda_graph_sample_kwarg_names", None + ) + assert ( + captured_sample_kwarg_names is None + or captured_sample_kwarg_names == sample_kwarg_names + ), ( + "HyperConnectionHybridLayer TE CUDA graph captures must use a stable " + "keyword-input signature." + ) + self._te_cuda_graph_sample_kwarg_names = sample_kwarg_names + self._rebuild_te_cuda_graph_packed_seq_params(kwargs) + if self._inner_is_partial_moe_capture(): + hidden_states = args[0] if args else kwargs["hidden_states"] + aggregated, h_res, h_post, residual = self.hyper_connection( + hidden_states, return_residual=True + ) + inner_kwargs = dict(kwargs) + inner_kwargs.pop("hidden_states", None) + inner_outputs = list( + self.inner_layer._te_cuda_graph_capture(aggregated, **inner_kwargs) + ) + return tuple(inner_outputs) + (h_post, h_res, residual) + + record_offload_events = ( + isinstance(self.inner_layer, TransformerLayer) + and getattr(self.inner_layer, "offload_module_in_cuda_graph", False) + ) + if record_offload_events: + if args: + hidden_states = self.inner_layer.off_interface.backward_record(args[0]) + args = (hidden_states,) + args[1:] + else: + hidden_states = self.inner_layer.off_interface.backward_record( + kwargs.pop("hidden_states") + ) + kwargs["hidden_states"] = hidden_states + hidden_states, context = self.forward(*args, **kwargs) + outputs = [hidden_states] + if context is not None: + outputs.append(context) + if record_offload_events: + self.inner_layer.off_interface.forward_record() + return tuple(outputs) + + def _te_cuda_graph_replay(self, *args, **kwargs): + """Replay a whole-wrapper graph or resume an eager partial-MoE tail.""" + self._flatten_te_cuda_graph_packed_seq_params(kwargs) + sample_kwarg_names = getattr( + self, "_te_cuda_graph_sample_kwarg_names", None + ) + assert sample_kwarg_names is not None, ( + "HyperConnectionHybridLayer TE CUDA graph replay requires the keyword " + "signature recorded during capture." + ) + for key in tuple(kwargs): + if key != "hidden_states" and key not in sample_kwarg_names: + kwargs.pop(key) + delayed_offload = ( + isinstance(self.inner_layer, TransformerLayer) + and getattr(self.inner_layer.config, "delay_offload_until_cuda_graph", False) + ) + if delayed_offload: + self.inner_layer.off_interface.enter_replay() + try: + outputs = list(super()._te_cuda_graph_replay(*args, **kwargs)) + if delayed_offload: + self.inner_layer.off_interface.flush_delayed_groups() + + if self._inner_is_partial_moe_capture(): + residual = outputs.pop() + h_res = outputs.pop() + h_post = outputs.pop() + _, mlp_output_with_bias = ( + self.inner_layer.resume_moe_experts_after_partial_cudagraph( + outputs + ) + ) + hidden_states = self.hyper_connection.fused_h_res_h_post_bda( + h_res, + residual, + h_post, + mlp_output_with_bias, + dropout_prob=self.inner_layer.hidden_dropout, + training=self.training, + fused=self.inner_layer.config.bias_dropout_fusion, + manager=None, + ) + if ( + self.config.fp32_residual_connection + and self.config.params_dtype is not None + and hidden_states.dtype != self.config.params_dtype + ): + hidden_states = hidden_states.to(self.config.params_dtype) + return hidden_states, None + + return outputs[0], None + finally: + if delayed_offload: + self.inner_layer.off_interface.exit_replay() + + def _get_te_cuda_graph_replay_args(self, *args, **kwargs): + """Use the wrapped TransformerLayer's replay-argument normalization.""" + if not isinstance(getattr(self, "inner_layer", None), TransformerLayer): + return super()._get_te_cuda_graph_replay_args(*args, **kwargs) + + missing = object() + previous_microbatch = getattr(self.inner_layer, "current_microbatch", missing) + self.inner_layer.current_microbatch = getattr(self, "current_microbatch", 0) + try: + return self.inner_layer._get_te_cuda_graph_replay_args(*args, **kwargs) + finally: + if previous_microbatch is missing: + del self.inner_layer.current_microbatch + else: + self.inner_layer.current_microbatch = previous_microbatch + + def _get_submodules_under_cudagraphs(self): + if self._inner_is_partial_moe_capture(): + return [ + self.hyper_connection + ] + self.inner_layer._get_submodules_under_cudagraphs() + return super()._get_submodules_under_cudagraphs() + + def mamba_state_shapes_per_request(self) -> Optional[Tuple[Tuple[int], Tuple[int]]]: + """Delegate Mamba inference state shape requests to the wrapped layer.""" + if not hasattr(self.inner_layer, 'mamba_state_shapes_per_request'): + return None + return self.inner_layer.mamba_state_shapes_per_request() + + def _call_inner_layer( + self, + hidden_states: Tensor, + attention_mask: Tensor, + inference_context: Optional[BaseInferenceContext], + rotary_pos_emb: Optional[Tensor], + sequence_len_offset: Optional[Tensor], + packed_seq_params: Optional[PackedSeqParams], + padding_mask: Optional[Tensor], + input_ids: Optional[Tensor] = None, + ) -> Tuple[Tensor, Optional[Tensor]]: + from megatron.core.transformer.cuda_graphs import is_graph_capturing + + inner = self.inner_layer.forward if is_graph_capturing() else self.inner_layer + if isinstance(self.inner_layer, TransformerLayer): + output = inner( + hidden_states=hidden_states, + attention_mask=attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + sequence_len_offset=sequence_len_offset, + packed_seq_params=packed_seq_params, + padding_mask=padding_mask, + input_ids=input_ids, + _called_from_hybrid_mhc_wrapper=True, + ) + else: + # Mamba-like layers only consume the common HybridStack arguments. + output = inner( + hidden_states=hidden_states, + attention_mask=attention_mask, + inference_context=inference_context, + packed_seq_params=packed_seq_params, + ) + + if isinstance(output, tuple): + context = output[1] if len(output) > 1 else None + return output[0], context + return output, None + + def _call_inner_transformer_layer_without_local_bda( + self, + hidden_states: Tensor, + attention_mask: Tensor, + inference_context: Optional[BaseInferenceContext], + rotary_pos_emb: Optional[Tensor], + sequence_len_offset: Optional[Tensor], + packed_seq_params: Optional[PackedSeqParams], + padding_mask: Optional[Tensor], + input_ids: Optional[Tensor] = None, + ) -> Optional[Tuple[Tuple[Tensor, Optional[Tensor]], Optional[Tensor], float, bool]]: + """Return a raw branch output for split Hybrid TransformerLayer instances. + + Hybrid layers are normally attention-only or MLP/MoE-only. For those + layers, bypass the inner layer's local residual/BDA and let the mHC BDA + own that operation directly. + """ + if not isinstance(self.inner_layer, TransformerLayer): + return None + + layer = self.inner_layer + if InferenceMode.is_active() and layer.config.inference_fuse_tp_communication: + return None + + has_attention = not isinstance(layer.self_attention, IdentityOp) + has_cross_attention = not isinstance(layer.cross_attention, IdentityOp) + has_mlp = not isinstance(layer.mlp, IdentityOp) + + if has_cross_attention or has_attention == has_mlp: + return None + + if has_attention: + output_with_bias, attn_norm_manager, residual = ( + layer._forward_self_attention_output_with_bias( + hidden_states=hidden_states, + attention_mask=attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + packed_seq_params=packed_seq_params, + sequence_len_offset=sequence_len_offset, + ) + ) + output_with_bias = layer._group_offload_output_with_bias( + output_with_bias, attn_norm_manager, forced_released_tensors=[residual] + ) + return output_with_bias, None, layer.hidden_dropout, layer.config.bias_dropout_fusion + + output_with_bias, residual = layer._forward_mlp_output_with_bias( + hidden_states, + inference_context=inference_context, + padding_mask=padding_mask, + packed_seq_params=packed_seq_params, + input_ids=input_ids, + ) + if layer.mlp_norm_manager is not None: + output_with_bias = layer._group_offload_output_with_bias( + output_with_bias, layer.mlp_norm_manager, forced_released_tensors=[residual] + ) + layer.mlp_norm_manager = None + return output_with_bias, None, layer.hidden_dropout, layer.config.bias_dropout_fusion + + def forward( + self, + hidden_states: Tensor, + attention_mask: Optional[Tensor] = None, + inference_context: Optional[BaseInferenceContext] = None, + rotary_pos_emb: Optional[Tensor] = None, + sequence_len_offset: Optional[Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + padding_mask: Optional[Tensor] = None, + mhc_recompute_manager=None, + input_ids: Optional[Tensor] = None, + ) -> Tuple[Tensor, Optional[Tensor]]: + """Run the wrapped hybrid layer through one layer-boundary mHC update.""" + if mhc_recompute_manager is None: + mhc_recompute_manager = getattr(self, '_mhc_recompute_manager', None) + aggregated, h_res, h_post, residual = self.hyper_connection( + hidden_states, + mhc_recompute_manager=mhc_recompute_manager, + return_residual=True, + ) + fast_path_result = self._call_inner_transformer_layer_without_local_bda( + aggregated, + attention_mask, + inference_context, + rotary_pos_emb, + sequence_len_offset, + packed_seq_params, + padding_mask, + input_ids, + ) + + if fast_path_result is None: + layer_output, context = self._call_inner_layer( + aggregated, + attention_mask, + inference_context, + rotary_pos_emb, + sequence_len_offset, + packed_seq_params, + padding_mask, + input_ids, + ) + if self.config.fp32_residual_connection and aggregated.dtype != layer_output.dtype: + aggregated = aggregated.to(layer_output.dtype) + layer_output_with_bias = (layer_output - aggregated, None) + dropout_prob = 0.0 + bias_dropout_fusion = False + else: + layer_output_with_bias, context, dropout_prob, bias_dropout_fusion = fast_path_result + + layer_output = layer_output_with_bias[0] + if layer_output.shape != aggregated.shape: + raise RuntimeError( + "HyperConnectionHybridLayer requires wrapped branches to preserve " + f"hidden-state shape. Got {tuple(layer_output.shape)} from wrapped branch " + f"vs {tuple(aggregated.shape)} input." + ) + + is_last_in_recompute_block = bool( + mhc_recompute_manager is not None + and getattr(mhc_recompute_manager, "is_last_layer_in_recompute_block", False) + ) + mhc_bda_manager = None if is_last_in_recompute_block else mhc_recompute_manager + hidden_states = self.hyper_connection.fused_h_res_h_post_bda( + h_res, + residual, + h_post, + layer_output_with_bias, + dropout_prob=dropout_prob, + training=self.training, + fused=bias_dropout_fusion, + manager=mhc_bda_manager, + ) + if ( + self.config.fp32_residual_connection + and self.config.params_dtype is not None + and hidden_states.dtype != self.config.params_dtype + ): + hidden_states = hidden_states.to(self.config.params_dtype) + return hidden_states, context + + class HybridStack(MegatronModule): """ Constructor for the HybridStack class. @@ -72,6 +541,7 @@ class HybridStack(MegatronModule): pg_collection (ProcessGroupCollection): the required model communication process groups to use. is_mtp_layer (bool, optional): whether this is an MTP layer. Defaults to False. + mtp_layer_number (int, optional): enclosing MTP depth for nested MoE metrics. """ def __init__( @@ -87,6 +557,7 @@ def __init__( dtype=None, pg_collection: ProcessGroupCollection = None, is_mtp_layer: bool = False, + mtp_layer_number: Optional[int] = None, name: str | None = None, ) -> None: """ @@ -98,6 +569,7 @@ def __init__( self.post_layer_norm = post_layer_norm self.post_process = post_process self.is_mtp_layer = is_mtp_layer + self.mtp_layer_number = mtp_layer_number assert pg_collection is not None, "pg_collection must be provided for HybridStack" @@ -108,6 +580,8 @@ def __init__( self.input_tensor = None self.pg_collection = pg_collection + self._mhc_block_end_plan: Optional[List[bool]] = None + assert layer_type_list is not None, ( "layer_type_list must be provided. It should be pre-computed from " "--hybrid-layer-pattern by HybridModel." @@ -156,6 +630,39 @@ def __init__( pp_layer_offset=pp_layer_offset, name=(name + f".layers.{i}") if name is not None else None, ) + elif layer_type == LayerSymbols.CSA: + layer = build_module( + submodules.csa_layer, + config=self.config, + layer_number=layer_number, + pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, + add_layer_offset=False, + pp_layer_offset=pp_layer_offset, + name=(name + f".layers.{i}") if name is not None else None, + ) + elif layer_type == LayerSymbols.HCA: + layer = build_module( + submodules.hca_layer, + config=self.config, + layer_number=layer_number, + pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, + add_layer_offset=False, + pp_layer_offset=pp_layer_offset, + name=(name + f".layers.{i}") if name is not None else None, + ) + elif layer_type == LayerSymbols.WINDOW: + layer = build_module( + submodules.window_layer, + config=self.config, + layer_number=layer_number, + pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, + add_layer_offset=False, + pp_layer_offset=pp_layer_offset, + name=(name + f".layers.{i}") if name is not None else None, + ) elif layer_type == LayerSymbols.MLP: layer = build_module( submodules.mlp_layer, @@ -171,6 +678,7 @@ def __init__( config=self.config, layer_number=layer_number, pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, add_layer_offset=False, name=(name + f".layers.{i}") if name is not None else None, ) @@ -186,6 +694,12 @@ def __init__( ) else: raise ValueError("unexpected layer_type") + if self.is_mtp_layer and self.mtp_layer_number is not None: + self._set_mtp_layer_number_for_moe_metrics( + layer, self.mtp_layer_number + ) + if self.config.enable_hyper_connections: + layer = HyperConnectionHybridLayer(config=self.config, layer=layer) self.layers.append(layer) if self.config.cuda_graph_impl == "local": @@ -202,6 +716,32 @@ def __init__( eps=self.config.layernorm_epsilon, ) + if ( + self.config.enable_hyper_connections + and self.post_process + and not self.is_mtp_layer + ): + hc_mult = self.config.num_residual_streams + hc_dim = self.config.hidden_size * hc_mult + self.hc_head_fn = mark_keep_in_fp32(nn.Parameter(torch.randn(hc_mult, hc_dim))) + self.hc_head_base = mark_keep_in_fp32(nn.Parameter(torch.zeros(hc_mult))) + self.hc_head_scale = mark_keep_in_fp32(nn.Parameter(torch.ones(1))) + nn.init.xavier_uniform_(self.hc_head_fn) + if self.config.sequence_parallel: + setattr(self.hc_head_fn, 'sequence_parallel', True) + setattr(self.hc_head_base, 'sequence_parallel', True) + setattr(self.hc_head_scale, 'sequence_parallel', True) + + @staticmethod + def _set_mtp_layer_number_for_moe_metrics( + layer: torch.nn.Module, mtp_layer_number: int + ) -> None: + """Propagate the enclosing MTP depth to nested MTP MoE routers.""" + for module in layer.modules(): + router = getattr(module, "router", None) + if router is not None and getattr(router, "is_mtp_layer", False): + router.mtp_layer_number = mtp_layer_number + def set_input_tensor(self, input_tensor: Tensor): """Set input tensor to be used instead of forward()'s input. @@ -222,6 +762,49 @@ def mamba_state_shapes_per_request(self) -> Optional[Tuple[Tuple[int], Tuple[int return layer.mamba_state_shapes_per_request() return None + def _compute_mhc_block_end_plan(self) -> List[bool]: + """Compute deterministic per-layer mHC recompute block boundaries.""" + num_layers = len(self.layers) + block_ends: List[bool] = [False] * num_layers + if num_layers == 0: + return block_ends + + layers_per_block = self.config.mhc_recompute_layer_num + for layer_idx in range(num_layers): + is_last_in_stack = layer_idx == num_layers - 1 + block_ends[layer_idx] = is_last_in_stack or ( + layers_per_block is not None and (layer_idx + 1) % layers_per_block == 0 + ) + return block_ends + + def _build_mhc_recompute_layer_plan( + self, use_mhc_recompute: bool + ) -> Tuple[List[Optional[CheckpointManager]], List[bool]]: + """Build single-use recompute managers for this forward pass.""" + num_layers = len(self.layers) + if not use_mhc_recompute or num_layers == 0: + return [None] * num_layers, [False] * num_layers + + if self._mhc_block_end_plan is None: + self._mhc_block_end_plan = self._compute_mhc_block_end_plan() + block_ends = self._mhc_block_end_plan + + layer_managers: List[Optional[CheckpointManager]] = [None] * num_layers + manager = CheckpointManager() + for layer_idx in range(num_layers): + layer_managers[layer_idx] = manager + if block_ends[layer_idx] and layer_idx != num_layers - 1: + manager = CheckpointManager() + return layer_managers, block_ends + + @staticmethod + def _finalize_mhc_recompute_layer( + manager: Optional[CheckpointManager], hidden_states: Tensor, is_block_end: bool + ) -> None: + """Finalize the current mHC recompute block when its last layer finishes.""" + if manager is not None and is_block_end: + manager.discard_all_outputs_and_register_unified_recompute(hidden_states) + def forward( self, hidden_states: Union[Tensor, WrappedTensor], @@ -232,6 +815,7 @@ def forward( inference_params: Optional[BaseInferenceContext] = None, packed_seq_params: Optional[PackedSeqParams] = None, padding_mask=None, + input_ids: Optional[Tensor] = None, ): """ Forward function of the HybridStack class. @@ -247,6 +831,8 @@ def forward( inference_context (BaseInferenceContext): the inference parameters. rotary_pos_emb (Tensor, optional): the rotary positional embeddings. Defaults to None. + input_ids (Tensor, optional): Token IDs forwarded to hash-routed + TransformerLayer instances. Defaults to None. Returns: Tensor: the output tensor. """ @@ -261,6 +847,15 @@ def forward( if isinstance(hidden_states, WrappedTensor): hidden_states = hidden_states.unwrap() + if ( + self.config.enable_hyper_connections + and self.pre_process + and not self.is_mtp_layer + ): + hidden_states = HyperConnectionModule.input_expand( + hidden_states, self.config.num_residual_streams + ) + if inference_context and inference_context.is_static_batching(): # NOTE(bnorick): match BaseInferenceContext attributes for # mamba_ssm.utils.generation.BaseInferenceContext, @@ -308,6 +903,14 @@ def get_inner_quant_context(config, layer_number): def get_inner_quant_context(config, layer_number): return nullcontext() + use_mhc_recompute = ( + self.training + and self.config.enable_hyper_connections + and self.config.recompute_granularity == 'selective' + and "mhc" in self.config.recompute_modules + ) + mhc_layer_managers, mhc_block_ends = self._build_mhc_recompute_layer_plan(use_mhc_recompute) + with outer_fp8_context: if self.config.recompute_granularity == 'full' and self.training: hidden_states = checkpointed_forward( @@ -320,17 +923,22 @@ def get_inner_quant_context(config, layer_number): attention_bias=None, packed_seq_params=packed_seq_params, padding_mask=padding_mask, + input_ids=input_ids, use_inner_quantization_context=(use_inner_fp8_context or use_fp4_context), ) else: - for layer in self.layers: + for layer_idx, layer in enumerate(self.layers): # Layers have 1-indexed layer numbers attribute. inner_quant_context = get_inner_quant_context( self.config, layer.layer_number - 1 ) + mhc_manager = mhc_layer_managers[layer_idx] + if mhc_manager is not None: + mhc_manager.is_last_layer_in_recompute_block = mhc_block_ends[layer_idx] + with inner_quant_context: - if isinstance(layer, TransformerLayer): - hidden_states, _ = layer( + if isinstance(layer, (TransformerLayer, HyperConnectionHybridLayer)): + layer_kwargs = dict( hidden_states=hidden_states, attention_mask=attention_mask, inference_context=inference_context, @@ -339,6 +947,13 @@ def get_inner_quant_context(config, layer_number): packed_seq_params=packed_seq_params, padding_mask=padding_mask, ) + if mhc_manager is not None and isinstance( + layer, HyperConnectionHybridLayer + ): + layer_kwargs["mhc_recompute_manager"] = mhc_manager + if input_ids is not None: + layer_kwargs['input_ids'] = input_ids + hidden_states, _ = layer(**layer_kwargs) else: # MambaLayer, Expert, or MLP hidden_states = layer( hidden_states=hidden_states, @@ -353,6 +968,29 @@ def get_inner_quant_context(config, layer_number): if isinstance(hidden_states, tuple): hidden_states = hidden_states[0] + self._finalize_mhc_recompute_layer( + manager=mhc_manager, + hidden_states=hidden_states, + is_block_end=mhc_block_ends[layer_idx], + ) + + mhc_multistream = None + if ( + self.config.enable_hyper_connections + and self.post_process + and not self.is_mtp_layer + ): + if (self.config.mtp_num_layers or 0) > 0: + mhc_multistream = hidden_states + hidden_states = learned_output_contract( + hidden_states, + self.hc_head_fn, + self.hc_head_base, + self.hc_head_scale, + self.config.num_residual_streams, + self.config.layernorm_epsilon, + ) + # Final layer norm. if self.post_process and self.post_layer_norm: hidden_states = self.final_norm(hidden_states) @@ -363,6 +1001,8 @@ def get_inner_quant_context(config, layer_number): inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True ) + if mhc_multistream is not None: + return hidden_states, mhc_multistream return hidden_states def sharded_state_dict( @@ -387,6 +1027,7 @@ def sharded_state_dict( dict: The sharded state dictionary for the current object. """ + sharded_offsets = sharded_offsets or () sharded_state_dict = {} layer_prefix = f'{prefix}layers.' @@ -421,6 +1062,20 @@ def sharded_state_dict( ) ) + local_state_dict: dict = {} + self._save_to_state_dict(local_state_dict, '', keep_vars=True) + if local_state_dict: + metadata = ensure_metadata_has_dp_cp_group(metadata) + sharded_state_dict.update( + make_sharded_tensors_for_checkpoint( + local_state_dict, + prefix, + sharded_offsets=sharded_offsets, + tp_group=self.tp_group, + dp_cp_group=metadata['dp_cp_group'], + ) + ) + return sharded_state_dict diff --git a/megatron/core/models/hybrid/hybrid_layer_allocation.py b/megatron/core/models/hybrid/hybrid_layer_allocation.py index 67103fe67f1..9d6b4da917f 100644 --- a/megatron/core/models/hybrid/hybrid_layer_allocation.py +++ b/megatron/core/models/hybrid/hybrid_layer_allocation.py @@ -18,11 +18,15 @@ class Symbols: GDN = 'G' ATTENTION = "*" DS_ATTENTION = "D" + CSA = "C" + HCA = "H" + WINDOW = "W" MLP = "-" MOE = 'E' PIPE = '|' MTP_SEPARATOR = "/" - VALID_LAYERS = {MAMBA, GDN, ATTENTION, DS_ATTENTION, MLP, MOE} + VALID_LAYERS = {MAMBA, GDN, ATTENTION, DS_ATTENTION, CSA, HCA, WINDOW, MLP, MOE} + MLA_ATTENTION = {DS_ATTENTION, CSA, HCA, WINDOW} @classmethod def name_sorted_valid_layer_symbols(cls) -> list[str]: @@ -173,10 +177,10 @@ def get_hybrid_layer_counts(pattern: str) -> Dict[str, int]: Examples: >>> get_hybrid_layer_counts("M*M*") - {'*': 2, 'G': 0, 'D': 0, 'M': 2, '-': 0, 'E': 0} + {'*': 2, 'C': 0, 'D': 0, 'G': 0, 'H': 0, 'M': 2, '-': 0, 'E': 0, 'W': 0} >>> get_hybrid_layer_counts("M-M-|M-M*-/MM/MM") - {'*': 1, 'G': 0, 'D': 0, 'M': 8, '-': 4, 'E': 0} + {'*': 1, 'C': 0, 'D': 0, 'G': 0, 'H': 0, 'M': 8, '-': 4, 'E': 0, 'W': 0} """ parsed = parse_hybrid_pattern(pattern) counts = {symbol: 0 for symbol in Symbols.name_sorted_valid_layer_symbols()} @@ -292,9 +296,11 @@ def _validate_pattern(pattern: str, pattern_name: str, allow_pipe: bool = False) f"Valid symbols are: {valid_chars}" ) - # Disallow Attention + MLA/DSA hybridity. - if Symbols.ATTENTION in pattern and Symbols.DS_ATTENTION in pattern: - raise ValueError("Not supported to have both Attention and MLA/DSA in one model") + # MLA variants may coexist, but standard attention cannot share a model with them. + if Symbols.ATTENTION in pattern and any(symbol in pattern for symbol in Symbols.MLA_ATTENTION): + raise ValueError( + "Not supported to have both Attention and MLA/DSA/CSA/HCA/Window in one model" + ) def validate_segment_layers(segment: str) -> List[str]: @@ -320,9 +326,13 @@ def validate_segment_layers(segment: str) -> List[str]: f"one of {Symbols.VALID_LAYERS}" ) - # Disallow Attention + MLA/DSA hybridity. - if Symbols.ATTENTION in segment and Symbols.DS_ATTENTION in segment: - raise ValueError("Not supported to have both Attention and MLA/DSA in one model") + # MLA variants may coexist, but standard attention cannot share a model with them. + if Symbols.ATTENTION in segment and any( + symbol in segment for symbol in Symbols.MLA_ATTENTION + ): + raise ValueError( + "Not supported to have both Attention and MLA/DSA/CSA/HCA/Window in one model" + ) return layer_type_list diff --git a/megatron/core/models/hybrid/hybrid_layer_specs.py b/megatron/core/models/hybrid/hybrid_layer_specs.py index e1624293b5a..12d25c26da5 100755 --- a/megatron/core/models/hybrid/hybrid_layer_specs.py +++ b/megatron/core/models/hybrid/hybrid_layer_specs.py @@ -1,5 +1,7 @@ # Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved. +from dataclasses import replace from functools import partial +from typing import Optional from megatron.core.extensions.transformer_engine import ( TEColumnParallelLinear, @@ -9,6 +11,7 @@ TENorm, TERowParallelLinear, ) +from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add from megatron.core.models.gpt.moe_module_specs import ( get_inference_optimized_moe_spec, @@ -36,6 +39,9 @@ DSAttention, DSAttentionSubmodules, ) +from megatron.core.transformer.experimental_attention_variant.dsv4_module_specs import ( + get_dsv4_hybrid_module_spec_for_backend, +) from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.mlp import MLP, MLPSubmodules from megatron.core.transformer.multi_latent_attention import ( @@ -77,7 +83,11 @@ submodules=MultiTokenPredictionLayerSubmodules( enorm=TENorm, hnorm=TENorm, + # Hybrid MTP selects the combined projection normally and + # per-stream projections when mHC is enabled. eh_proj=TEColumnParallelLinear, + e_proj=TEColumnParallelLinear, + h_proj=TEColumnParallelLinear, mtp_model_layer=None, # Built via pattern + hybrid_submodules layer_norm=TENorm, ), @@ -296,7 +306,10 @@ submodules=MultiTokenPredictionLayerSubmodules( enorm=TENorm, hnorm=TENorm, + # Keep both projection forms available for Hybrid MTP. eh_proj=InferenceColumnParallelLinear, + e_proj=InferenceColumnParallelLinear, + h_proj=InferenceColumnParallelLinear, mtp_model_layer=None, # Built via pattern + hybrid_submodules layer_norm=TENorm, ), @@ -311,3 +324,42 @@ # Backward-compatible aliases mamba_stack_spec = hybrid_stack_spec mamba_inference_stack_spec = hybrid_inference_stack_spec + + +def hybrid_dsv4_stack_spec(config: TransformerConfig) -> ModuleSpec: + """Build a HybridStack whose D/C/H/W symbols use DeepSeek-V4 attention. + + D reads its compression ratio from ``config.csa_compress_ratios``. C, H, and W + encode fixed ratios 4, 128, and 0 respectively. + """ + assert config.transformer_impl == "transformer_engine", ( + "DSv4 HybridModel currently supports only the transformer-engine implementation." + ) + dsv4_attention = get_dsv4_hybrid_module_spec_for_backend( + config=config, backend=TESpecProvider() + ) + + def wrap_dsv4_layer(compress_ratio: Optional[int] = None) -> ModuleSpec: + attention = dsv4_attention + if compress_ratio is not None: + attention = replace( + dsv4_attention, + params={**dsv4_attention.params, "compress_ratio": compress_ratio}, + ) + return ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + input_layernorm=TENorm, + self_attention=attention, + self_attn_bda=get_bias_dropout_add, + ), + ) + + submodules = replace( + hybrid_stack_spec.submodules, + dsa_layer=wrap_dsv4_layer(), + csa_layer=wrap_dsv4_layer(compress_ratio=4), + hca_layer=wrap_dsv4_layer(compress_ratio=128), + window_layer=wrap_dsv4_layer(compress_ratio=0), + ) + return ModuleSpec(module=HybridStack, submodules=submodules) diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index f750c77e05b..258c658f384 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -511,15 +511,25 @@ def forward( # be None, so this assert will succeed. # assert attention_mask is None, "The attention mask is ignored and should be set to None" + decoder_extra_block_kwargs = {} + if self.config.moe_n_hash_layers > 0 and input_ids is not None: + decoder_extra_block_kwargs['input_ids'] = input_ids + # Run decoder. - hidden_states = self.decoder( + decoder_output = self.decoder( hidden_states=decoder_input, attention_mask=attention_mask, inference_context=inference_context, rotary_pos_emb=rotary_pos_emb, packed_seq_params=packed_seq_params, padding_mask=padding_mask, + **decoder_extra_block_kwargs, ) + if isinstance(decoder_output, tuple): + hidden_states, mhc_multistream = decoder_output + else: + hidden_states = decoder_output + mhc_multistream = None output_weight = None if self.share_embeddings_and_output_weights: @@ -541,11 +551,13 @@ def forward( input_ids=input_ids, position_ids=position_ids, hidden_states=hidden_states, + mhc_multistream=mhc_multistream, attention_mask=attention_mask, inference_params=inference_params, rotary_pos_emb=rotary_pos_emb, packed_seq_params=packed_seq_params, embedding=self.embedding, + padding_mask=padding_mask, ) if not self.post_process: diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index e503a16dde3..d7e2618f3f7 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -7,7 +7,6 @@ import math import warnings from abc import ABC, abstractmethod -from itertools import chain from logging import getLogger from typing import Any, Callable, Dict, List, Optional, Tuple, Union @@ -1019,14 +1018,42 @@ def sharded_state_dict( state_dict = self.state_dict() + # Optimizer state ids enumerate the inner optimizer params: the fp32 main + # copies of float16 params and the native fp32 params, interleaved in the + # original param-group order. Yield the model-side param for each inner + # param in that order so the ids line up even when both kinds are present. + def model_params_in_optimizer_order(): + for inner_group, float16_group, fp32_group in zip( + self.optimizer.param_groups, self.float16_groups, self.fp32_from_fp32_groups + ): + float16_params = iter(float16_group) + fp32_param_ids = {id(param) for param in fp32_group} + for param in inner_group['params']: + yield param if id(param) in fp32_param_ids else next(float16_params) + id_to_sharded_param_map = get_param_id_to_sharded_param_map( - model_sharded_state_dict, chain.from_iterable(g for g in self.float16_groups) + model_sharded_state_dict, model_params_in_optimizer_order() ) # Convert fp32_from_fp16_params assert len(state_dict['fp32_from_fp16_params']) == len( state_dict['optimizer']['param_groups'] ) + # State ids of the fp32 main copies only, skipping native fp32 params. + float16_param_ids_per_group = [] + for state_group, inner_group, fp32_group in zip( + state_dict['optimizer']['param_groups'], + self.optimizer.param_groups, + self.fp32_from_fp32_groups, + ): + fp32_param_ids = {id(param) for param in fp32_group} + float16_param_ids_per_group.append( + [ + param_id + for param_id, param in zip(state_group['params'], inner_group['params']) + if id(param) not in fp32_param_ids + ] + ) state_dict['fp32_from_fp16_params'] = [ [ make_sharded_optimizer_tensor( @@ -1034,10 +1061,10 @@ def sharded_state_dict( fp32_param, prefix=f'optimizer.state.fp32_param', ) - for param_id, fp32_param in zip(state_group['params'], fp32_group) + for param_id, fp32_param in zip(param_ids, fp32_group) ] - for fp32_group, state_group in zip( - state_dict['fp32_from_fp16_params'], state_dict['optimizer']['param_groups'] + for fp32_group, param_ids in zip( + state_dict['fp32_from_fp16_params'], float16_param_ids_per_group ) ] diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index 1c26af56244..7dd985cccb9 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -1,10 +1,31 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. -from dataclasses import dataclass +from dataclasses import dataclass, replace +from typing import Literal, Mapping, MutableMapping, Optional, Tuple, Union import torch import torch.distributed as dist +import torch.nn.functional as F from torch import Tensor +CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX = "_packed_seq_params_" + +PACKED_SEQ_PARAMS_CUDA_GRAPH_TENSOR_FIELDS = ( + "cu_seqlens_q", + "cu_seqlens_kv", + "cu_seqlens_q_padded", + "cu_seqlens_kv_padded", +) + +PACKED_SEQ_PARAMS_CUDA_GRAPH_STATIC_FIELDS = ( + "qkv_format", + "max_seqlen_q", + "max_seqlen_kv", + "local_cp_size", + "cp_group", + "pad_between_seqs", + "cp_partition_mode", +) + @dataclass class PackedSeqParams: @@ -26,6 +47,7 @@ class PackedSeqParams: seq_idx: Tensor = None tokens_per_sample: int = None pad_between_seqs: bool = None + cp_partition_mode: Literal["zigzag", "contiguous"] = "zigzag" def __post_init__(self): """Pre-compute seq_idx for Mamba mixer CUDA graph compatibility. @@ -66,3 +88,421 @@ def __post_init__(self): .to(torch.int32) .unsqueeze(0) # Add a batch dimension ) + + +def _cuda_graph_packed_seq_params_key(field_name: str, prefix: str) -> str: + return f"{prefix}{field_name}" + + +def split_packed_seq_params_for_cuda_graph( + packed_seq_params: PackedSeqParams | None, prefix: str = CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX +) -> tuple[dict[str, Tensor | None], dict[str, object]]: + """Split ``PackedSeqParams`` into graph Tensor inputs and static metadata. + + Transformer Engine CUDA graph inputs must be tensors or ``None``. ``PackedSeqParams`` mixes + dynamic Tensor fields, such as cumulative sequence lengths, with static metadata, such as THD + format and max sequence lengths. This helper keeps only the fields TE attention consumes; + Mamba-only fields such as ``total_tokens`` and ``seq_idx`` stay outside this graph boundary. + """ + if packed_seq_params is None: + return {}, {} + + tensor_kwargs = {} + for field_name in PACKED_SEQ_PARAMS_CUDA_GRAPH_TENSOR_FIELDS: + value = getattr(packed_seq_params, field_name) + if value is not None and not isinstance(value, Tensor): + raise TypeError( + f"PackedSeqParams.{field_name} must be a Tensor or None for CUDA graphs, " + f"got {type(value).__name__}." + ) + if value is not None: + tensor_kwargs[_cuda_graph_packed_seq_params_key(field_name, prefix)] = value + + static_metadata = {} + for field_name in PACKED_SEQ_PARAMS_CUDA_GRAPH_STATIC_FIELDS: + value = getattr(packed_seq_params, field_name) + if isinstance(value, Tensor): + raise TypeError( + f"PackedSeqParams.{field_name} is static CUDA graph metadata and must not be " + "a Tensor." + ) + static_metadata[field_name] = value + + return tensor_kwargs, static_metadata + + +def has_packed_seq_params_cuda_graph_kwargs( + kwargs: Mapping[str, object], prefix: str = CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX +) -> bool: + """Return whether ``kwargs`` contains flattened ``PackedSeqParams`` Tensor fields.""" + return any( + _cuda_graph_packed_seq_params_key(field_name, prefix) in kwargs + for field_name in PACKED_SEQ_PARAMS_CUDA_GRAPH_TENSOR_FIELDS + ) + + +def build_packed_seq_params_from_cuda_graph_kwargs( + kwargs: MutableMapping[str, object], + static_metadata: Mapping[str, object] | None, + prefix: str = CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX, + remove_from_kwargs: bool = True, +) -> PackedSeqParams | None: + """Rebuild ``PackedSeqParams`` from flattened CUDA graph kwargs. + + Args: + kwargs: Graph kwargs that may contain flattened packed-sequence Tensor fields. + static_metadata: Non-Tensor metadata produced by + :func:`split_packed_seq_params_for_cuda_graph`. + prefix: Prefix used for flattened Tensor fields. + remove_from_kwargs: Whether to pop consumed flattened fields from ``kwargs``. + """ + packed_seq_params_kwargs = dict(static_metadata or {}) + found_tensor_field = False + for field_name in PACKED_SEQ_PARAMS_CUDA_GRAPH_TENSOR_FIELDS: + key = _cuda_graph_packed_seq_params_key(field_name, prefix) + if key not in kwargs: + continue + found_tensor_field = True + value = kwargs.pop(key) if remove_from_kwargs else kwargs[key] + if value is not None and not isinstance(value, Tensor): + raise TypeError( + f"Flattened PackedSeqParams field {key} must be a Tensor or None, " + f"got {type(value).__name__}." + ) + packed_seq_params_kwargs[field_name] = value + + if not packed_seq_params_kwargs and not found_tensor_field: + return None + + return PackedSeqParams(**packed_seq_params_kwargs) + + +def _pad_seq_tensor(tensor: Optional[Tensor], target_len: int) -> Optional[Tensor]: + """Pad a token-like tensor along its last dimension with zeros.""" + if tensor is None: + return None + actual_len = tensor.shape[-1] + assert actual_len <= target_len, ( + f"Sequence-length tensor (last dim = {actual_len}) exceeds target ({target_len}); " + "increase max_seqlen_per_dp_cp_rank or filter overlong samples upstream." + ) + if actual_len == target_len: + return tensor + return F.pad(tensor, (0, target_len - actual_len), value=0) + + +def _pad_padding_mask(mask: Tensor, target_len: int) -> Tensor: + """Pad a boolean padding mask with ``True`` along its last dimension.""" + actual_len = mask.shape[-1] + assert actual_len <= target_len, ( + f"Padding mask length ({actual_len}) exceeds target ({target_len}); " + "refusing to silently truncate." + ) + if actual_len == target_len: + return mask + + pad_shape = list(mask.shape) + pad_shape[-1] = target_len - actual_len + tail = torch.ones(pad_shape, dtype=mask.dtype, device=mask.device) + return torch.cat((mask, tail), dim=-1) + + +def _pad_cu_seqlens(cu_seqlens: Optional[Tensor], target_entries: int) -> Optional[Tensor]: + """Pad cumulative-length metadata to a static number of entries.""" + if cu_seqlens is None: + return None + actual_entries = cu_seqlens.shape[0] + assert actual_entries <= target_entries, ( + f"Actual num_seqs ({actual_entries - 1}) exceeds thd_max_packed_sequences " + f"({target_entries - 1})." + ) + if actual_entries == target_entries: + return cu_seqlens + padded = torch.empty( + (target_entries,), dtype=cu_seqlens.dtype, device=cu_seqlens.device + ) + padded.fill_(cu_seqlens[-1].item()) + padded[:actual_entries] = cu_seqlens + return padded + + +def _append_dummy_seq(cu_seqlens: Optional[Tensor], dummy_end: int) -> Optional[Tensor]: + """Append a cumulative boundary for the post-pack padding tail.""" + if cu_seqlens is None: + return None + dummy = torch.full( + (1,), int(dummy_end), dtype=cu_seqlens.dtype, device=cu_seqlens.device + ) + return torch.cat((cu_seqlens, dummy), dim=0) + + +def _round_up_to_alignment(value: int, alignment: int) -> int: + assert alignment > 0, f"Packed sequence padding alignment must be > 0, got {alignment}." + return ((value + alignment - 1) // alignment) * alignment + + +def get_thd_padding_kwargs( + pad_packed_seq_alignment: Union[int, Literal["max"]], + max_seqlen_per_dp_cp_rank: Optional[int], + thd_max_packed_sequences: Optional[int], + cuda_graph_static: bool, +) -> Tuple[Optional[int], Optional[int], Optional[int]]: + """Resolve token and cumulative-length padding settings from the config.""" + if cuda_graph_static: + assert max_seqlen_per_dp_cp_rank is not None + return None, int(max_seqlen_per_dp_cp_rank), thd_max_packed_sequences + + if pad_packed_seq_alignment == "max": + assert max_seqlen_per_dp_cp_rank is not None + return None, int(max_seqlen_per_dp_cp_rank), None + + return int(pad_packed_seq_alignment), None, None + + +def _resolve_thd_cp_geometry( + packed_seq_params: PackedSeqParams, + cp_group: Optional[dist.ProcessGroup] = None, + cp_size: Optional[int] = None, + cp_rank: Optional[int] = None, +) -> Tuple[int, int]: + """Resolve THD context-parallel geometry from explicitly threaded state.""" + if cp_group is not None: + return int(dist.get_world_size(group=cp_group)), int(dist.get_rank(group=cp_group)) + + if cp_size is None: + if packed_seq_params.cp_group is not None: + return ( + int(dist.get_world_size(group=packed_seq_params.cp_group)), + int(dist.get_rank(group=packed_seq_params.cp_group)), + ) + cp_size = packed_seq_params.local_cp_size or 1 + + cp_size = int(cp_size) + if cp_size == 1: + return 1, 0 + if cp_rank is None: + raise ValueError( + "cp_rank or cp_group must be provided when padding THD metadata with cp_size > 1." + ) + return cp_size, int(cp_rank) + + +def _resolve_thd_padding_lengths( + tokens: Optional[Tensor], + labels: Optional[Tensor], + loss_mask: Optional[Tensor], + position_ids: Optional[Tensor], + packed_seq_params: PackedSeqParams, + target_len: Optional[int], + alignment: Optional[int], + cp_group: Optional[dist.ProcessGroup] = None, + cp_size: Optional[int] = None, + cp_rank: Optional[int] = None, + padding_mask: Optional[Tensor] = None, +) -> Tuple[int, int, int, int, torch.device]: + """Resolve local and global THD padding lengths without changing tensors.""" + cp_size, cp_rank = _resolve_thd_cp_geometry( + packed_seq_params, cp_group=cp_group, cp_size=cp_size, cp_rank=cp_rank + ) + + local_tensor_len = None + mask_device = None + for candidate in (tokens, labels, loss_mask, position_ids, padding_mask): + if candidate is not None: + local_tensor_len = int(candidate.shape[-1]) + mask_device = candidate.device + break + + has_local_tensor = local_tensor_len is not None + if packed_seq_params.cu_seqlens_q is not None: + global_actual_len = int(packed_seq_params.cu_seqlens_q[-1].item()) + if mask_device is None: + mask_device = packed_seq_params.cu_seqlens_q.device + else: + assert has_local_tensor, ( + "packed_seq_params.cu_seqlens_q must be available to derive padding_mask " + "when all token-like tensors are None." + ) + global_actual_len = local_tensor_len * cp_size + + if has_local_tensor: + local_actual_len = local_tensor_len + local_target_len = ( + int(target_len) + if target_len is not None + else _round_up_to_alignment(local_actual_len, alignment) + ) + return ( + local_actual_len, + global_actual_len, + local_target_len, + local_target_len * cp_size, + mask_device, + ) + + global_target_len = ( + int(target_len) * cp_size + if target_len is not None + else _round_up_to_alignment(global_actual_len, alignment) + ) + if cp_size > 1: + from megatron.core.extensions.transformer_engine import get_thd_partitioned_indices + + partition_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 + ) + local_actual_len = int( + get_thd_partitioned_indices( + partition_cu_seqlens, global_actual_len, cp_size, cp_rank + ).numel() + ) + local_target_len = int( + get_thd_partitioned_indices( + partition_cu_seqlens, global_target_len, cp_size, cp_rank + ).numel() + ) + else: + local_actual_len = global_actual_len + local_target_len = global_target_len + + return ( + local_actual_len, + global_actual_len, + local_target_len, + global_target_len, + mask_device, + ) + + +def pad_sequence_for_thd( + tokens: Optional[Tensor], + labels: Optional[Tensor], + loss_mask: Optional[Tensor], + position_ids: Optional[Tensor], + packed_seq_params: PackedSeqParams, + alignment: Optional[int] = None, + target_len: Optional[int] = None, + max_num_seqs: Optional[int] = None, + pad_by_appending_dummy_seq: bool = True, + padding_mask: Optional[Tensor] = None, + cp_group: Optional[dist.ProcessGroup] = None, + cp_size: Optional[int] = None, + cp_rank: Optional[int] = None, +) -> Tuple[ + Optional[Tensor], + Optional[Tensor], + Optional[Tensor], + Optional[Tensor], + PackedSeqParams, + Optional[Tensor], +]: + """Pad packed THD tensors and return a mask for the padding tail.""" + assert (alignment is None) != ( + target_len is None + ), "Exactly one of alignment or target_len must be provided for THD padding." + + ( + local_actual_len, + global_actual_len, + local_target_len, + global_target_len, + mask_device, + ) = _resolve_thd_padding_lengths( + tokens, + labels, + loss_mask, + position_ids, + packed_seq_params, + target_len=target_len, + alignment=alignment, + cp_group=cp_group, + cp_size=cp_size, + cp_rank=cp_rank, + padding_mask=padding_mask, + ) + + if packed_seq_params.cu_seqlens_q is not None: + cu_seqlens = packed_seq_params.cu_seqlens_q + individual_lens = cu_seqlens[1:] - cu_seqlens[:-1] + max_individual = ( + int(individual_lens.max().item()) if individual_lens.numel() > 0 else 0 + ) + assert max_individual <= global_target_len, ( + f"Individual request length ({max_individual}) exceeds the global padded " + f"capacity ({global_target_len})." + ) + + tokens = _pad_seq_tensor(tokens, local_target_len) + labels = _pad_seq_tensor(labels, local_target_len) + loss_mask = _pad_seq_tensor(loss_mask, local_target_len) + position_ids = _pad_seq_tensor(position_ids, local_target_len) + + cu_seqlens_q = packed_seq_params.cu_seqlens_q + cu_seqlens_kv = packed_seq_params.cu_seqlens_kv + cu_seqlens_q_padded = packed_seq_params.cu_seqlens_q_padded + cu_seqlens_kv_padded = packed_seq_params.cu_seqlens_kv_padded + + target_cu_entries = None if max_num_seqs is None else max_num_seqs + 1 + has_dummy_padding_seq = ( + pad_by_appending_dummy_seq and global_target_len > global_actual_len + ) + dummy_seq_len = ( + global_target_len - global_actual_len if has_dummy_padding_seq else 0 + ) + if has_dummy_padding_seq: + cu_seqlens_q = _append_dummy_seq(cu_seqlens_q, global_target_len) + cu_seqlens_kv = _append_dummy_seq(cu_seqlens_kv, global_target_len) + cu_seqlens_q_padded = _append_dummy_seq( + cu_seqlens_q_padded, global_target_len + ) + cu_seqlens_kv_padded = _append_dummy_seq( + cu_seqlens_kv_padded, global_target_len + ) + + if target_cu_entries is not None: + cu_seqlens_q = _pad_cu_seqlens(cu_seqlens_q, target_cu_entries) + cu_seqlens_kv = _pad_cu_seqlens(cu_seqlens_kv, target_cu_entries) + cu_seqlens_q_padded = _pad_cu_seqlens( + cu_seqlens_q_padded, target_cu_entries + ) + cu_seqlens_kv_padded = _pad_cu_seqlens( + cu_seqlens_kv_padded, target_cu_entries + ) + + padded_params = replace( + packed_seq_params, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + cu_seqlens_q_padded=cu_seqlens_q_padded, + cu_seqlens_kv_padded=cu_seqlens_kv_padded, + max_seqlen_q=( + global_target_len + if target_cu_entries is not None + else max(packed_seq_params.max_seqlen_q or 0, dummy_seq_len) + ), + max_seqlen_kv=( + global_target_len + if target_cu_entries is not None + else max(packed_seq_params.max_seqlen_kv or 0, dummy_seq_len) + ), + total_tokens=local_target_len if target_cu_entries is None else None, + seq_idx=None, + pad_between_seqs=( + False if has_dummy_padding_seq else packed_seq_params.pad_between_seqs + ), + ) + + tail_padding_mask = ( + torch.arange(local_target_len, device=mask_device).unsqueeze(0) + >= local_actual_len + ) + if padding_mask is None: + padding_mask = tail_padding_mask + else: + padding_mask = ( + _pad_padding_mask(padding_mask, local_target_len) | tail_padding_mask + ) + + return tokens, labels, loss_mask, position_ids, padded_params, padding_mask diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py index e67c498e2cc..20b11a126d4 100644 --- a/megatron/core/pipeline_parallel/schedules.py +++ b/megatron/core/pipeline_parallel/schedules.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import contextlib from functools import partial @@ -277,7 +277,7 @@ def _get_experimental_attention_variant_loss_scale_func(config): if loss_scale_func is not None: return loss_scale_func - if getattr(config, 'experimental_attention_variant', None) == 'dsa': + if getattr(config, 'experimental_attention_variant', None) in ('dsa', 'dsv4_hybrid'): from megatron.core.transformer.experimental_attention_variant.dsa import ( DSAIndexerLossAutoScaler, ) @@ -1161,7 +1161,11 @@ def enable_grad_sync(): model_type = get_model_type(model[0]) - tensor_shape = [seq_length, micro_batch_size, config.hidden_size] + tensor_shape = [ + seq_length, + micro_batch_size, + _get_pipeline_hidden_size(config, pp_group=p2p_communicator.pp_group), + ] tensor_shape[0] = tensor_shape[0] // cp_group.size() if config.sequence_parallel: tensor_shape[0] = tensor_shape[0] // tp_group.size() @@ -2100,11 +2104,17 @@ def get_tensor_shapes( config, tp_group: Optional[torch.distributed.ProcessGroup] = None, cp_group: Optional[torch.distributed.ProcessGroup] = None, + pp_group: Optional[torch.distributed.ProcessGroup] = None, + is_recv: bool = True, ): """Determine tensor shapes for pipeline communication. Returns [()] for variable_seq_lengths mode (shapes exchanged dynamically), or computed shapes for fixed sequence length mode. + + ``pp_group`` and ``is_recv`` distinguish the two mHC boundary shapes. The + first stage has no forward receive and the last stage has no forward send; + all communication between stages carries the expanded residual streams. """ tensor_shapes = [] @@ -2120,10 +2130,35 @@ def get_tensor_shapes( if config.sequence_parallel: effective_seq_length = effective_seq_length // tp_group.size() - tensor_shapes.append((effective_seq_length, micro_batch_size, config.hidden_size)) + hidden_size = _get_pipeline_hidden_size(config, pp_group=pp_group, is_recv=is_recv) + tensor_shapes.append((effective_seq_length, micro_batch_size, hidden_size)) return tensor_shapes +def _get_pipeline_hidden_size( + config, *, pp_group: Optional[torch.distributed.ProcessGroup], is_recv: Optional[bool] = None +) -> int: + """Return the hidden dimension used by a pipeline communication edge.""" + hidden_size = config.hidden_size + if not getattr(config, 'enable_hyper_connections', False) or pp_group is None: + return hidden_size + + pp_rank = pp_group.rank() + pp_size = pp_group.size() + if pp_size == 1: + return hidden_size + + # The interleaved schedule uses one shape for all P2P operations. Its only + # active communication edges are between stages, so all of them carry n*C. + if is_recv is None: + return hidden_size * config.num_residual_streams + + is_inactive_boundary = (is_recv and pp_rank == 0) or (not is_recv and pp_rank == pp_size - 1) + if is_inactive_boundary: + return hidden_size + return hidden_size * config.num_residual_streams + + def forward_backward_pipelining_without_interleaving( *, forward_step_func, @@ -2287,6 +2322,9 @@ def enable_grad_sync(): else: backward_func = backward_step + # Multi-module pipelines exchange variable shapes and do not expose one + # pipeline group for the full topology. + pp_group = getattr(p2p_communicator, 'pp_group', None) recv_tensor_shapes = get_tensor_shapes( seq_length=seq_length, micro_batch_size=micro_batch_size, @@ -2294,6 +2332,8 @@ def enable_grad_sync(): config=config, tp_group=tp_group, cp_group=cp_group, + pp_group=pp_group, + is_recv=True, ) send_tensor_shapes = get_tensor_shapes( seq_length=seq_length, @@ -2302,6 +2342,8 @@ def enable_grad_sync(): config=config, tp_group=tp_group, cp_group=cp_group, + pp_group=pp_group, + is_recv=False, ) if adjust_tensor_shapes_fn is not None: recv_tensor_shapes, send_tensor_shapes = adjust_tensor_shapes_fn( diff --git a/megatron/core/recompute.py b/megatron/core/recompute.py index bd0d1bcb3b2..75dbd8717d9 100644 --- a/megatron/core/recompute.py +++ b/megatron/core/recompute.py @@ -31,6 +31,7 @@ def checkpointed_forward( padding_mask: Optional[Tensor] = None, extract_layer_indices: Optional[Set[int]] = None, layer_offset: int = 0, + input_ids: Optional[Tensor] = None, ) -> Union[Tensor, Tuple[Tensor, Tensor]]: """Forward method with activation checkpointing. @@ -41,6 +42,7 @@ def checkpointed_forward( layer_offset (int): The global layer offset for the current pipeline stage. Used to convert local layer indices to global indices when checking extract_layer_indices. + input_ids (Tensor, optional): Token IDs forwarded to hash-routed MoE layers. Returns: If extract_layer_indices is empty: hidden_states tensor @@ -64,6 +66,7 @@ def custom_forward( rotary_pos_emb_local, rotary_pos_emb_global, padding_mask=None, + input_ids=None, ): rotary_pos_emb = ( (rotary_pos_emb_local, rotary_pos_emb_global) @@ -106,11 +109,19 @@ def custom_forward( packed_seq_params=packed_seq_params, padding_mask=padding_mask, ) + if input_ids is not None: + layer_kwargs["input_ids"] = input_ids with inner_quantization_context: if isinstance(layer, TransformerLayer): hidden_states, context = layer(**layer_kwargs) else: # MambaLayer (HybridStack `M` slot) - for k in ("context", "context_mask", "attention_bias", "padding_mask"): + for k in ( + "context", + "context_mask", + "attention_bias", + "padding_mask", + "input_ids", + ): layer_kwargs.pop(k, None) hidden_states = layer(**layer_kwargs) context = None @@ -126,7 +137,15 @@ def chunk_runner(start: int, end: int, use_checkpoint: bool): nonlocal hidden_states, context cf = custom(start, end) # Unpack the RoPE tuple as torch cannot save tuples for backward pass. - args = (hidden_states, attention_mask, context, context_mask, *rotary_pos_emb, padding_mask) + args = ( + hidden_states, + attention_mask, + context, + context_mask, + *rotary_pos_emb, + padding_mask, + input_ids, + ) if use_checkpoint: # Precision-aware activation checkpoint: TE under FP8/FP4, # tensor_parallel under BF16/FP16/FP32. diff --git a/megatron/core/tensor_parallel/random.py b/megatron/core/tensor_parallel/random.py index 4cf945dd8bb..fb79ec82004 100644 --- a/megatron/core/tensor_parallel/random.py +++ b/megatron/core/tensor_parallel/random.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Parts of the code here are adapted from PyTorch # repo: https://github.com/pytorch/pytorch @@ -598,7 +598,9 @@ def forward( @staticmethod def backward(ctx, *args): """Backward pass.""" - if not torch.autograd._is_checkpoint_valid(): + from megatron.core.transformer.cuda_graphs import is_graph_capturing + + if not torch.autograd._is_checkpoint_valid() and not is_graph_capturing(): raise RuntimeError( "Checkpointing is not compatible with .grad(), " "please use .backward() if possible" @@ -649,10 +651,67 @@ def checkpoint( return CheckpointFunction.apply(function, distribute_saved_activations, *args) +def _save_args_to_ctx(ctx, args): + """Save mixed tensor/non-tensor arguments into autograd ctx. + + Since save_for_backward only supports tensors, this function separates + tensor and non-tensor arguments, saving tensors via save_for_backward + and storing non-tensor metadata (indices and values) as ctx attributes. + + Use _load_args_from_ctx to reconstruct the original args. + """ + tensor_args = [] + non_tensor_entries = [] + + for index, arg in enumerate(args): + if isinstance(arg, torch.Tensor): + tensor_args.append(arg) + continue + non_tensor_entries.append((index, arg)) + + ctx.save_for_backward(*detach_variable(tuple(tensor_args))) + ctx._non_tensor_entries = tuple(non_tensor_entries) + ctx._total_args_count = len(args) + + +def _load_args_from_ctx(ctx): + """Load and reconstruct mixed tensor/non-tensor arguments from autograd ctx. + + This is the inverse of _save_args_to_ctx. It retrieves tensors from + ctx.saved_tensors and merges them with stored non-tensor arguments + to reconstruct the original args in their original order. + + Returns: + tuple of reconstructed arguments in their original order. + """ + + def _detach_with_grad(tensor): + detached = tensor.detach() + detached.requires_grad_(tensor.requires_grad) + return detached + + tensor_iter = iter(_detach_with_grad(t) for t in ctx.saved_tensors) + total_args_count = ctx._total_args_count + non_tensor_map = dict(ctx._non_tensor_entries) + + reconstructed_args = [] + for index in range(total_args_count): + if index in non_tensor_map: + reconstructed_args.append(non_tensor_map[index]) + else: + reconstructed_args.append(next(tensor_iter)) + return tuple(reconstructed_args) + + class CheckpointWithoutOutputFunction(torch.autograd.Function): """ Checkpoint Function Helper for CheckpointWithoutOutput. Save context for recompute. + + Handles both tensor and non-tensor arguments: + - Tensor arguments are saved via save_for_backward + - Non-tensor arguments (int, float, bool, None, etc.) are stored separately + in ctx attributes and reconstructed during recomputation """ @staticmethod @@ -675,7 +734,10 @@ def forward( with torch.no_grad(), fwd_ctx: outputs = run_function(*args) - ctx.save_for_backward(*detach_variable(args)) + + # Save tensor and non-tensor arguments into ctx for recomputation + _save_args_to_ctx(ctx, args) + # the CheckpointWithoutOutput object is passed in, then it can access the saved input # tensors later for recomputation checkpoint_without_output_obj.ctx = ctx @@ -692,10 +754,60 @@ def backward(ctx, *args): torch.autograd.backward(outputs, args) ctx.outputs = None ctx.inputs = None - grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else inp for inp in inputs) + grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else None for inp in inputs) return (None, None) + grads +class CheckpointManager: + """ + Manages multiple CheckpointWithoutOutput objects within a TransformerBlock + cross layer recomputations, enabling unified recomputation during backward pass. + This is particularly useful for scenarios where multiple checkpoint operations have + sequential dependencies (i.e., the output of one checkpoint is the input of the next). + + Usage: + ckptManager = CheckpointManager() + ckpt_function = CheckpointWithoutOutput(ckpt_manager=ckptManager) + ckpt_function.checkpoint(run_function, *args) + # other checkpointed operations + ckpt_manager.discard_all_outputs_and_register_unified_recompute(final_output) + """ + + def __init__(self): + self.checkpoints = [] + # Set by TransformerBlock before each layer forward. + # When True, the layer should keep block-boundary output uncheckpointed. + self.is_last_layer_in_recompute_block = False + + def add_checkpoint(self, ckpt): + """Add a checkpoint to the manager.""" + if not isinstance(ckpt, CheckpointWithoutOutput): + raise TypeError("Expected CheckpointWithoutOutput object") + if ckpt.outputs is None: + raise ValueError("CheckpointWithoutOutput must call checkpoint() before adding") + self.checkpoints.append(ckpt) + + def discard_all_outputs_and_register_unified_recompute(self, hook_tensor): + """Discard all checkpoint outputs to save memory and register unified recompute hook.""" + for ckpt in self.checkpoints: + for output in ckpt.outputs: + output.untyped_storage().resize_(0) + + # Register unified recompute hook + if hook_tensor.requires_grad: + hook_tensor.register_hook(self._unified_recompute_hook) + + def _unified_recompute_hook(self, grad_output): + for ckpt in self.checkpoints: + # Call _recompute for each checkpoint in forward order + # The _recompute method will restore the output tensor storage + ckpt._recompute(None) + + +# Compatibility for the already-reviewed mHC prerequisite API. +CheckpointWithoutOutputManager = CheckpointManager + + class CheckpointWithoutOutput(object): """ Checkpoint a model or part of the model and release the output. @@ -710,8 +822,19 @@ class CheckpointWithoutOutput(object): discarded output tensors are directly saved in the following modules for backward computation. """ - def __init__(self, fp8=False): - self.fp8 = fp8 is not None + def __init__(self, fp8=False, ckpt_manager=None): + """ + Initialize CheckpointWithoutOutput. + + Args: + fp8: Whether to use FP8 mode. Defaults to False. + ckpt_manager: Optional CheckpointManager instance. When provided, + checkpoint() will auto-register to the manager, and + discard_output_and_register_recompute() will only discard + output without registering individual hooks. + """ + self.fp8 = bool(fp8) + self.ckpt_manager = ckpt_manager self.run_function = None self.fwd_cpu_rng_state = None self.fwd_cuda_rng_state = None @@ -720,7 +843,12 @@ def __init__(self, fp8=False): self.outputs = None def checkpoint(self, run_function: Callable[[Unpack[_Ts]], _R], *args: Unpack[_Ts]) -> _R: - """Checkpoint function.""" + """ + Checkpoint function. + + If ckpt_manager was provided during initialization, this checkpoint + will be automatically registered to the manager after execution. + """ # If in cuda graph warmup, disable checkpointing, as 'discard_output_and_register_recompute' # may be called in a separate graph warmup. @@ -737,6 +865,11 @@ def checkpoint(self, run_function: Callable[[Unpack[_Ts]], _R], *args: Unpack[_T self.outputs = outputs if isinstance(self.outputs, torch.Tensor): self.outputs = (self.outputs,) + + # Auto-register to manager if provided + if self.ckpt_manager is not None: + self.ckpt_manager.add_checkpoint(self) + return outputs def _recompute(self, _): @@ -745,7 +878,7 @@ def _recompute(self, _): from megatron.core.transformer.cuda_graphs import is_graph_capturing, is_graph_warmup # The recomputation has been triggered already. Just return. - # Handle cudagraphs, do nothing if currently in graph warmup + # Handle cudagraphs: do nothing if currently in graph warmup if self.ctx is None or is_graph_warmup(): return @@ -767,17 +900,8 @@ def _recompute(self, _): recompute_ctx = contextlib.nullcontext() fp8_ctx = contextlib.nullcontext() - # Store the inputs for backward pass - inputs = self.ctx.saved_tensors - - def detach(t): - if isinstance(t, torch.Tensor): - requires_grad = t.requires_grad - t = t.detach() - t.requires_grad_(requires_grad) - return t - - inputs = tuple(detach(t) for t in inputs) + # Reconstruct full args list from saved ctx + inputs = _load_args_from_ctx(self.ctx) with torch.enable_grad(), fp8_ctx, recompute_ctx: outputs = self.run_function(*inputs) @@ -810,10 +934,11 @@ def discard_output_and_register_recompute(self, hook_tensor): in the forward pass and the gradient of the hook_tensor is computed before the recomputed tensors are used. """ - + # When ckpt_manager is set, this is a no-op. + # Manager handles all discarding and hook registration uniformly. from megatron.core.transformer.cuda_graphs import is_graph_warmup - if is_graph_warmup(): + if self.ckpt_manager is not None or is_graph_warmup(): return # use resize to release the output tensor memory and still keep the metadata in the tensors. diff --git a/megatron/core/transformer/__init__.py b/megatron/core/transformer/__init__.py index 0e3cdcfa57e..75e3b485c4f 100644 --- a/megatron/core/transformer/__init__.py +++ b/megatron/core/transformer/__init__.py @@ -1,6 +1,10 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from .module import MegatronModule from .spec_utils import ModuleSpec, build_module from .transformer_config import MLATransformerConfig, TransformerConfig -from .transformer_layer import TransformerLayer, TransformerLayerSubmodules +from .transformer_layer import ( + HyperConnectionTransformerLayer, + TransformerLayer, + TransformerLayerSubmodules, +) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 7b1a09ea333..412550b5d24 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -303,6 +303,7 @@ def __init__( cp_comm_type: str | None = None, pg_collection: ProcessGroupCollection | None = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, name: str | None = None, ): """ @@ -314,6 +315,7 @@ def __init__( self.config = config self.layer_number = layer_number self._pp_layer_offset = pp_layer_offset + self.is_mtp_layer = is_mtp_layer self.attn_mask_type = attn_mask_type self.attention_type = attention_type @@ -1463,8 +1465,11 @@ def forward( cu_seqlens_kv = packed_seq_params.cu_seqlens_kv_padded else: cu_seqlens_kv = packed_seq_params.cu_seqlens_kv + rope_max_seqlen_q = packed_seq_params.max_seqlen_q + rope_max_seqlen_kv = packed_seq_params.max_seqlen_kv else: cu_seqlens_q = cu_seqlens_kv = None + rope_max_seqlen_q = rope_max_seqlen_kv = None if split_qkv: if q_pos_emb is not None: @@ -1477,6 +1482,7 @@ def forward( cu_seqlens=cu_seqlens_q, mscale=self._yarn_concentration_factor, cp_group=self.pg_collection.cp, + max_seqlen=rope_max_seqlen_q, ) else: query = inference_context.apply_rotary_emb_query( @@ -1495,6 +1501,7 @@ def forward( cu_seqlens=cu_seqlens_kv, mscale=self._yarn_concentration_factor, cp_group=self.pg_collection.cp, + max_seqlen=rope_max_seqlen_kv, ) else: query, key, value = apply_fused_qkv_rotary_pos_emb( @@ -1631,6 +1638,7 @@ def __init__( cp_comm_type: str | None = None, pg_collection: ProcessGroupCollection | None = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, name: str | None = None, ): """ @@ -1646,6 +1654,7 @@ def __init__( cp_comm_type=cp_comm_type, pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, name=name, ) @@ -2047,6 +2056,7 @@ def __init__( attn_mask_type: AttnMaskType = AttnMaskType.padding, cp_comm_type: str | None = None, pg_collection: ProcessGroupCollection | None = None, + is_mtp_layer: bool = False, name: str | None = None, ): """ @@ -2061,6 +2071,7 @@ def __init__( attention_type="cross", cp_comm_type=cp_comm_type, pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, name=name, ) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 210f39fa217..bf80da99b23 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -21,6 +21,7 @@ from torch.utils._pytree import tree_map as tree_map_pyt from megatron.core.num_microbatches_calculator import get_num_microbatches +from megatron.core.packed_seq_params import split_packed_seq_params_for_cuda_graph from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.random import ( CudaRNGStatesTracker, @@ -1634,24 +1635,39 @@ def _layer_is_graphable(layer, config): if not isinstance(layer, GraphableMegatronModule): return False - # If cuda_graph_modules is not set, every layer is graphed. - if not config.cuda_graph_modules: - return True - # import modules here to avoid a circular import + from megatron.core.models.hybrid.hybrid_block import HyperConnectionHybridLayer from megatron.core.ssm.mamba_layer import MambaLayer from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.mlp import MLP from megatron.core.transformer.moe.moe_layer import MoELayer from megatron.core.transformer.transformer_layer import TransformerLayer - if isinstance(layer, MambaLayer) and CudaGraphModule.mamba in config.cuda_graph_modules: + inspected_layer = layer.inner_layer if isinstance(layer, HyperConnectionHybridLayer) else layer + + if ( + isinstance(layer, HyperConnectionHybridLayer) + and isinstance(inspected_layer, MambaLayer) + and getattr(config, "sequence_packing_scheduler", None) is not None + ): + # The mHC wrapper does not yet define a static packed-sequence + # metadata contract for its inner Mamba layer. + return False + + # If cuda_graph_modules is not set, every remaining layer is graphed. + if not config.cuda_graph_modules: + return True + + if ( + isinstance(inspected_layer, MambaLayer) + and CudaGraphModule.mamba in config.cuda_graph_modules + ): # mamba layer. return True - if isinstance(layer, TransformerLayer): + if isinstance(inspected_layer, TransformerLayer): if CudaGraphModule.attn in config.cuda_graph_modules and not ( - isinstance(layer.self_attention, IdentityOp) - and isinstance(layer.cross_attention, IdentityOp) + isinstance(inspected_layer.self_attention, IdentityOp) + and isinstance(inspected_layer.cross_attention, IdentityOp) ): # attn layer. return True @@ -1659,15 +1675,42 @@ def _layer_is_graphable(layer, config): CudaGraphModule.moe in config.cuda_graph_modules or CudaGraphModule.moe_router in config.cuda_graph_modules or CudaGraphModule.moe_preprocess in config.cuda_graph_modules - ) and isinstance(layer.mlp, MoELayer): + ) and isinstance(inspected_layer.mlp, MoELayer): # moe layer. return True - if CudaGraphModule.mlp in config.cuda_graph_modules and isinstance(layer.mlp, MLP): + if CudaGraphModule.mlp in config.cuda_graph_modules and isinstance( + inspected_layer.mlp, MLP + ): # mlp layer. return True return False +def _add_packed_seq_params_to_te_cuda_graph_sample_kwargs( + layer, sample_kwargs, sample_packed_seq_params +): + """Add flattened ``PackedSeqParams`` Tensor inputs to TE graph sample kwargs.""" + if sample_packed_seq_params is None: + return + + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph( + sample_packed_seq_params + ) + duplicate_keys = set(sample_kwargs) & set(tensor_kwargs) + assert not duplicate_keys, ( + "PackedSeqParams CUDA graph Tensor kwargs overlap with existing sample kwargs: " + f"{', '.join(sorted(duplicate_keys))}." + ) + assert hasattr(layer, '_set_te_cuda_graph_packed_seq_params_static_metadata'), ( + "Transformer layers using TE CUDA graph packed sequence samples must support " + "PackedSeqParams static metadata." + ) + layer._set_te_cuda_graph_packed_seq_params_static_metadata( + static_metadata, tensor_kwargs.keys() + ) + sample_kwargs.update(tensor_kwargs) + + class TECudaGraphHelper: """ Helper class to capture CUDA Graphs using TE make_graphed_callables(). @@ -1678,7 +1721,15 @@ class TECudaGraphHelper: """ def __init__( - self, model, config, seq_length, micro_batch_size, optimizers=[], pg_collection=None + self, + model, + config, + seq_length, + micro_batch_size, + optimizers=[], + pg_collection=None, + sample_packed_seq_params=None, + thd_sequence_length_upper_bound=None, ): assert HAVE_TE_GRAPHS, "CUDA Graphs are not supported without TE." assert ( @@ -1692,14 +1743,21 @@ def __init__( "CUDA Graph with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True." ) self.model = model + assert sample_packed_seq_params is None or is_te_min_version("1.10.0"), ( + "TE CUDA graph packed_seq_params support requires Transformer Engine >= 1.10.0 " + "because packed-sequence Tensor fields are passed as keyword arguments." + ) self.config = config self.seq_length = seq_length + self.thd_sequence_length_upper_bound = thd_sequence_length_upper_bound self.micro_batch_size = micro_batch_size self.optimizers = optimizers self.pg_collection = pg_collection + self.sample_packed_seq_params = sample_packed_seq_params if self.pg_collection is None: self.pg_collection = ProcessGroupCollection.use_mpu_process_groups() self.tp_group = self.pg_collection.tp + self.dp_group = self.pg_collection.dp self.dp_cp_group = self.pg_collection.dp_cp self.pp_group = self.pg_collection.pp from megatron.core.pipeline_parallel.p2p_communication import P2PCommunicator @@ -1904,12 +1962,16 @@ def get_rotary_pos_emb(transformer_module, transformer_input): static_inputs = layer.get_layer_static_inputs(self.seq_length, self.micro_batch_size) + from megatron.core.models.hybrid.hybrid_block import HyperConnectionHybridLayer from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.transformer_layer import TransformerLayer + inspected_layer = ( + layer.inner_layer if isinstance(layer, HyperConnectionHybridLayer) else layer + ) contains_self_attn = ( - isinstance(layer, TransformerLayer) - and not isinstance(layer.self_attention, IdentityOp) + isinstance(inspected_layer, TransformerLayer) + and not isinstance(inspected_layer.self_attention, IdentityOp) and ( not self.config.cuda_graph_modules or CudaGraphModule.attn in self.config.cuda_graph_modules @@ -1925,6 +1987,21 @@ def get_rotary_pos_emb(transformer_module, transformer_input): rotary_pos_emb = get_rotary_pos_emb(chunk_of_the_layer, hidden_states) if rotary_pos_emb is not None: static_inputs["rotary_pos_emb"] = rotary_pos_emb + static_packed_seq_params = static_inputs.pop( + "packed_seq_params", None + ) + assert not ( + static_packed_seq_params is not None + and self.sample_packed_seq_params is not None + ), ( + "Packed-sequence graph capture received both layer-generated static " + "metadata and an explicit sample_packed_seq_params." + ) + _add_packed_seq_params_to_te_cuda_graph_sample_kwargs( + layer, + static_inputs, + static_packed_seq_params or self.sample_packed_seq_params, + ) _sample_kwargs = static_inputs elif contains_self_attn: _sample_args = ( @@ -2075,6 +2152,120 @@ def _get_amax_reduction_group(self, with_context_parallel=False, tp_only_amax_re assert self.pg_collection.tp is not None return self.pg_collection.tp + def _should_use_dynamic_microbatch_slots(self) -> bool: + """Whether variable packed batches need a bounded graph-slot pool.""" + return bool(getattr(self.config, "cuda_graph_dynamic_microbatches", False)) + + @staticmethod + def _get_required_num_microbatch_slots_from_order(order, num_model_chunks): + """Infer the peak per-chunk forward liveness from a PP/VPP order.""" + outstanding = [0] * num_model_chunks + max_outstanding = [0] * num_model_chunks + for chunk_id in order: + if ceil(chunk_id) != chunk_id: + continue + model_chunk_idx = abs(int(ceil(chunk_id))) - 1 + if chunk_id > 0: + outstanding[model_chunk_idx] += 1 + max_outstanding[model_chunk_idx] = max( + max_outstanding[model_chunk_idx], + outstanding[model_chunk_idx], + ) + else: + outstanding[model_chunk_idx] -= 1 + assert outstanding[model_chunk_idx] >= 0, ( + "Invalid PP/VPP schedule while inferring CUDA graph slots." + ) + assert all(count == 0 for count in outstanding), ( + "Invalid PP/VPP schedule: outstanding forwards did not drain." + ) + return max(1, max(max_outstanding, default=1)) + + def _get_probe_num_microbatches_for_dynamic_slots(self): + """Return a topology-only microbatch count for slot-liveness probing.""" + pipeline_parallel_size = self.pp_group.size() + if pipeline_parallel_size == 1 and not self.config.overlap_moe_expert_parallel_comm: + return 1 + group_size = ( + self.config.microbatch_group_size_per_vp_stage + or pipeline_parallel_size + ) + return max( + pipeline_parallel_size * max(1, self.num_model_chunks) * 4, + group_size * max(1, self.num_model_chunks) * 2, + 1, + ) + + @staticmethod + def _get_dp_balanced_thd_max_num_microbatches( + global_batch_size, + dp_size, + cp_size, + max_seqlen_per_dp_cp_rank, + max_sequence_length, + microbatch_group_size_per_vp_stage=None, + max_num_seqs=None, + ): + """Return a conservative packed-microbatch upper bound.""" + assert global_batch_size >= 1 + assert dp_size >= 1 + assert cp_size >= 1 + assert max_seqlen_per_dp_cp_rank >= 1 + assert max_sequence_length >= 1 + + packed_capacity = max_seqlen_per_dp_cp_rank * cp_size + sequences_per_pack = max(1, packed_capacity // max_sequence_length) + if max_num_seqs is not None: + sequences_per_pack = min(sequences_per_pack, max(1, int(max_num_seqs))) + + num_packs = math.ceil(global_batch_size / sequences_per_pack) + multiple = dp_size * (microbatch_group_size_per_vp_stage or 1) + num_packs = math.ceil(num_packs / multiple) * multiple + return max(1, num_packs // dp_size) + + def _get_thd_varlen_max_num_microbatches( + self, runtime_num_microbatches, microbatch_group_size_per_vp_stage + ): + """Return the dp-balanced THD upper bound used for graph capture.""" + if ( + self.config.sequence_packing_scheduler != 'dp_balanced' + or self.config.max_seqlen_per_dp_cp_rank is None + ): + return runtime_num_microbatches, "runtime" + + dp_size = self.dp_group.size() + cp_size = self.dp_cp_group.size() // dp_size + global_batch_size = runtime_num_microbatches * self.micro_batch_size * dp_size + max_sequence_length = ( + self.thd_sequence_length_upper_bound + if self.thd_sequence_length_upper_bound is not None + else self.seq_length + ) + max_num_seqs = getattr(self.config, 'thd_max_packed_sequences', None) + if max_num_seqs is not None: + max_num_seqs = int(max_num_seqs) + if getattr(self.config, 'pad_packed_seq_alignment', None) is not None and getattr( + self.config, 'pad_packed_seq_by_appending_dummy_seq', True + ): + max_num_seqs -= 1 + + return ( + self._get_dp_balanced_thd_max_num_microbatches( + global_batch_size, + dp_size, + cp_size, + int(self.config.max_seqlen_per_dp_cp_rank), + int(max_sequence_length), + microbatch_group_size_per_vp_stage=( + None + if self.config.virtual_pipeline_model_parallel_size is None + else microbatch_group_size_per_vp_stage + ), + max_num_seqs=max_num_seqs, + ), + "thd_varlen_upper_bound", + ) + def _get_cuda_graph_input_data(self): """ Create the CUDA Graph capturing input data. @@ -2087,26 +2278,99 @@ def _get_cuda_graph_input_data(self): get_schedule_table, ) + microbatch_group_size_per_vp_stage = ( + self.config.microbatch_group_size_per_vp_stage or self.pp_group.size() + ) + # If PP is not enabled, we only need to capture one microbatch. if self.pp_group.size() == 1 and not self.config.overlap_moe_expert_parallel_comm: assert ( self.num_model_chunks == 1 ), "If PP is not enabled, there should be only one model chunk." self.num_microbatches = 1 + elif self._should_use_dynamic_microbatch_slots(): + probe_num_microbatches = self._get_probe_num_microbatches_for_dynamic_slots() + _, _, probe_warmup_microbatches, _ = get_pp_rank_microbatches( + probe_num_microbatches, + self.num_model_chunks, + microbatch_group_size_per_vp_stage, + False, + overlap_moe_expert_parallel_comm=( + self.config.overlap_moe_expert_parallel_comm + ), + ) + probe_schedule = get_schedule_table( + probe_num_microbatches, + self.num_model_chunks, + microbatch_group_size_per_vp_stage, + ) + probe_order = convert_schedule_table_to_order( + probe_warmup_microbatches, + self.num_model_chunks, + probe_schedule, + ) + auto_num_slots = self._get_required_num_microbatch_slots_from_order( + probe_order, self.num_model_chunks + ) + if self.pp_group.size() > 1: + auto_num_slots_tensor = torch.tensor( + [auto_num_slots], + dtype=torch.int32, + device=torch.cuda.current_device(), + ) + torch.distributed.all_reduce( + auto_num_slots_tensor, + op=torch.distributed.ReduceOp.MAX, + group=self.pp_group, + ) + auto_num_slots = int(auto_num_slots_tensor.item()) + + runtime_num_microbatches = get_num_microbatches() + max_num_microbatches, capture_mode = ( + self._get_thd_varlen_max_num_microbatches( + runtime_num_microbatches, + microbatch_group_size_per_vp_stage, + ) + ) + if self.config.overlap_moe_expert_parallel_comm or self.config.delay_wgrad_compute: + self.num_microbatches = runtime_num_microbatches + capture_mode = "runtime" + fallback_reason = "overlap_moe_expert_parallel_comm/delay_wgrad_compute" + else: + # The topology-only liveness value is a useful lower bound, but + # TE currently requires capture against the full THD/GBS-derived + # upper bound when the real packed-microbatch count can change. + self.num_microbatches = max( + runtime_num_microbatches, max_num_microbatches + ) + fallback_reason = None + log_on_each_pipeline_stage( + logger=logger, + tp_group=self.tp_group, + dp_cp_group=self.dp_cp_group, + level=logging.INFO, + msg=f'Rank {torch.distributed.get_rank()}: dynamic CUDA graph slots ' + f'enabled. runtime_num_microbatches={runtime_num_microbatches}, ' + f'auto_num_slots={auto_num_slots}, ' + f'max_num_microbatches={max_num_microbatches}, ' + f'capture_num_microbatches={self.num_microbatches}, ' + f'capture_mode={capture_mode}' + + (f', fallback_reason={fallback_reason}' if fallback_reason else ''), + ) else: self.num_microbatches = get_num_microbatches() _, _, num_warmup_microbatches, _ = get_pp_rank_microbatches( self.num_microbatches, self.num_model_chunks, - self.config.microbatch_group_size_per_vp_stage, + microbatch_group_size_per_vp_stage, forward_only=False, p2p_communicator=self.p2p_communicator, ) schedule_table = get_schedule_table( self.num_microbatches, self.num_model_chunks, - self.config.microbatch_group_size_per_vp_stage, + microbatch_group_size_per_vp_stage, ) order = convert_schedule_table_to_order( num_warmup_microbatches, self.num_model_chunks, schedule_table diff --git a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py index fccf674d785..cb03e3c5c94 100644 --- a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py +++ b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py @@ -146,6 +146,7 @@ def __init__( pg_collection: ProcessGroupCollection = None, pp_layer_offset: Optional[int] = None, name: str | None = None, + is_mtp_layer: bool = False, ): if pg_collection is None: pg_collection = ProcessGroupCollection.use_mpu_process_groups() @@ -160,6 +161,7 @@ def __init__( pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, name=name, + is_mtp_layer=is_mtp_layer, ) assert not config.add_bias_linear, "add_bias_linear is not supported for AbsorbedMLA" @@ -442,8 +444,11 @@ def get_query_key_value_tensors( cu_seqlens_kv = packed_seq_params.cu_seqlens_kv_padded else: cu_seqlens_kv = packed_seq_params.cu_seqlens_kv + rope_max_seqlen_q = packed_seq_params.max_seqlen_q + rope_max_seqlen_kv = packed_seq_params.max_seqlen_kv else: cu_seqlens_q = cu_seqlens_kv = None + rope_max_seqlen_q = rope_max_seqlen_kv = None # ========================================= # Q down projection @@ -631,6 +636,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po mscale=mscale, cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, + max_seqlen=rope_max_seqlen_q, ) # k_pos_emb:[num_tokens, 1, qk_pos_emb_head_dim] k_pos_emb = apply_rotary_pos_emb( @@ -641,6 +647,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po mscale=mscale, cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, + max_seqlen=rope_max_seqlen_kv, ) # query: [num_tokens, n, (kv_lora_rank + qk_pos_emb_head_dim)] diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py new file mode 100644 index 00000000000..4fe85b05ba0 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -0,0 +1,2737 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import copy +from dataclasses import dataclass +from functools import lru_cache +from typing import Optional, Tuple, Union + +import torch +import torch.nn as nn + +from megatron.core.fp8_utils import get_fp8_disabled_context +from megatron.core.fusions.fused_mla_yarn_rope_apply import fused_mla_rope_inplace +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.csa_kernels import ( + FusedIndexerSparseAttnFromTopkFunc, + batch_of_row, + build_flat_topk_idxs, + dsa_sparse_attn, + fused_indexer_sparse_attn, + indexer_topk, +) +from megatron.core.transformer.experimental_attention_variant.dsa import ( + DSAIndexerLossAutoScaler, + DSAIndexerLossLoggingHelper, + FusedDSAIndexerLoss, + fused_qk_topk_naive, + fused_qk_topk_naive_thd, + rotate_activation, +) +from megatron.core.transformer.module import MegatronModule, mark_keep_in_fp32 +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import nvtx_range_pop, nvtx_range_push + +# --------------------------------------------------------------------------- +# Helper functions for index computation +# --------------------------------------------------------------------------- + + +@lru_cache(maxsize=8) +def _get_window_topk_idxs_cached(window_size: int, seqlen: int, device_str: str) -> torch.Tensor: + """Compute sliding-window indices for a single sequence (cached). + + Returns: + indices: [seqlen, window_size] int tensor, -1 for invalid positions. + """ + base = torch.arange(seqlen, device=device_str).unsqueeze(1) + offsets = torch.arange(window_size, device=device_str) + matrix = (base - window_size + 1).clamp(min=0) + offsets + matrix = torch.where(matrix > base, -1, matrix) + return matrix + + +def get_window_topk_idxs( + window_size: int, batch_size: int, seqlen: int, device: torch.device +) -> torch.Tensor: + """Sliding-window indices [batch, seqlen, window_size].""" + matrix = _get_window_topk_idxs_cached(window_size, seqlen, str(device)) + return matrix.unsqueeze(0).expand(batch_size, -1, -1) + + +@lru_cache(maxsize=8) +def _get_compress_topk_idxs_cached( + ratio: int, seqlen: int, offset: int, device_str: str +) -> torch.Tensor: + """Compute all-compressed-positions indices for a single sequence (cached). + + Returns: + indices: [seqlen, seqlen // ratio] int tensor, -1 for future positions. + """ + n_compressed = seqlen // ratio + matrix = torch.arange(n_compressed, device=device_str).repeat(seqlen, 1) + mask = matrix >= torch.arange(1, seqlen + 1, device=device_str).unsqueeze(1) // ratio + matrix = torch.where(mask, -1, matrix + offset) + return matrix + + +def get_compress_topk_idxs( + ratio: int, batch_size: int, seqlen: int, offset: int, device: torch.device +) -> torch.Tensor: + """All-compressed-position indices [batch, seqlen, seqlen // ratio].""" + matrix = _get_compress_topk_idxs_cached(ratio, seqlen, offset, str(device)) + return matrix.unsqueeze(0).expand(batch_size, -1, -1) + + +def _get_csa_compressed_capacity( + packed_seq_params: Optional[PackedSeqParams], ratio: int, total_tokens: int +) -> Optional[int]: + """Return a host-known compressed capacity for THD CUDA graph capture. + + The exact ``sum(seq_len // ratio)`` lives in ``cu_seqlens`` on device. + Reading it on the host would break CUDA graph capture, and + ``PackedSeqParams`` must stay a generic MCore contract rather than + carrying CSA-specific metadata. Use a static upper bound instead; + device-side ``cu_seqlens_compressed`` keeps the true valid rows and + downstream kernels leave the extra rows as tail padding. + """ + if packed_seq_params is None or ratio <= 1: + return None + max_seqlen = packed_seq_params.max_seqlen_q + cu_seqlens = ( + packed_seq_params.cu_seqlens_q_padded + if packed_seq_params.cu_seqlens_q_padded is not None + else packed_seq_params.cu_seqlens_q + ) + if max_seqlen is None or cu_seqlens is None: + return None + num_sequences = max(int(cu_seqlens.shape[0]) - 1, 0) + return min(int(total_tokens) // ratio, num_sequences * (int(max_seqlen) // ratio)) + + +# --------------------------------------------------------------------------- +# THD (packed) variants of the index helpers above. +# +# Both produce per-row local-to-segment indices in the SAME index space that +# ``csa_kernels.local_to_global_flat(..., cu_seqlens_q=..., cu_seqlens_kv=...)`` +# expects: each row is one query token in the packed layout, each value is +# either ``-1`` (invalid / future position) or a non-negative local KV id in +# ``[0, seqlen_kv_full[batch_of_row])`` where ``seqlen_kv_full[b] = +# seqlen_kv[b] + seqlen_compressed[b]``. Window indices live in +# ``[0, seqlen_kv[b])``; compressed indices live in +# ``[seqlen_kv[b], seqlen_kv[b] + seqlen_compressed[b])``. +# +# These mirror the SBHD helpers above but cannot be lru-cached because +# their output shape depends on the per-batch ``cu_seqlens`` tensors. +# --------------------------------------------------------------------------- + + +def get_window_topk_idxs_thd( + window_size: int, cu_seqlens_q: torch.Tensor, total_q: Optional[int] = None +) -> torch.Tensor: + """Sliding-window indices for a packed THD layout. + + For each query token ``i`` in segment ``b`` (with ``pos_in_seq = + i - cu_seqlens_q[b]``), the window covers the last ``window_size`` + KV positions within the same segment's original KV region: + indices ``[max(0, pos-window_size+1), ..., pos]``; positions + extending before the start of the segment are emitted as ``-1``. + + Args: + window_size: number of positions per window. + cu_seqlens_q: ``(B+1,)`` int32 cumulative Q lengths + (self-attention: same as KV lengths). + total_q: total number of query tokens (avoids a GPU→CPU sync + when the caller already knows it, e.g. from ``x.shape[0]``). + + Returns: + ``(total_q, window_size)`` int32 — LOCAL (per-segment) KV indices. + """ + if total_q is None: + total_q = int(cu_seqlens_q[-1].item()) + device = cu_seqlens_q.device + batch_of_token = batch_of_row(cu_seqlens_q, total_q=total_q) + token_idx = torch.arange(total_q, device=device, dtype=cu_seqlens_q.dtype) + valid = token_idx < cu_seqlens_q[-1] + pos_in_seq = token_idx - cu_seqlens_q[batch_of_token] + pos_in_seq = torch.where(valid, pos_in_seq, torch.zeros_like(pos_in_seq)) + + offsets = torch.arange(window_size, device=device, dtype=cu_seqlens_q.dtype) + matrix = (pos_in_seq - window_size + 1).clamp(min=0).unsqueeze(1) + offsets.unsqueeze(0) + matrix = torch.where(matrix > pos_in_seq.unsqueeze(1), torch.full_like(matrix, -1), matrix) + matrix = torch.where(valid.unsqueeze(1), matrix, torch.full_like(matrix, -1)) + return matrix.int() + + +def get_compress_topk_idxs_thd( + ratio: int, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, + total_q: Optional[int] = None, + max_n_compressed: Optional[int] = None, +) -> torch.Tensor: + """All compressed-position indices for a packed THD layout. + + For each query token ``i`` in segment ``b`` (``pos_in_seq = i - + cu_seqlens_q[b]``), the valid compressed positions within that + segment are ``[0, 1, ..., (pos+1) // ratio - 1]`` (clamped to + ``seqlen_compressed[b]``). The returned indices are already shifted + by the per-segment offset ``seqlen_kv[b]`` so that they live in the + *full* per-segment KV index space ``[seqlen_kv[b], seqlen_kv[b] + + seqlen_compressed[b])`` — exactly mirroring the SBHD helper's + ``offset=sq`` shift. + + Args: + ratio: indexer compression ratio. + cu_seqlens_q: ``(B+1,)`` int32 cumulative Q lengths. + cu_seqlens_kv: ``(B+1,)`` int32 cumulative original-KV lengths + (used to derive the per-segment compressed-offset). + cu_seqlens_compressed: ``(B+1,)`` int32 cumulative compressed-KV + lengths (== Compressor's second return value). + total_q: total number of query tokens (avoids a GPU→CPU sync + when the caller already knows it, e.g. from ``x.shape[0]``). + max_n_compressed: max compressed sequence length across segments + (avoids a GPU→CPU sync when the caller can derive it, e.g. + ``max_seqlen_q // ratio``). + + Returns: + ``(total_q, max_compressed_per_seq)`` int32 — LOCAL (per-segment) + full-KV indices, ``-1`` for future positions. + """ + if total_q is None: + total_q = int(cu_seqlens_q[-1].item()) + device = cu_seqlens_q.device + seq_lens_kv = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] + seq_lens_compressed = cu_seqlens_compressed[1:] - cu_seqlens_compressed[:-1] + if max_n_compressed is None: + if seq_lens_compressed.numel() == 0: + return torch.empty((total_q, 0), dtype=torch.int32, device=device) + max_n_compressed = int(seq_lens_compressed.max().item()) + if max_n_compressed == 0: + return torch.empty((total_q, 0), dtype=torch.int32, device=device) + + batch_of_token = batch_of_row(cu_seqlens_q, total_q=total_q) + token_idx = torch.arange(total_q, device=device, dtype=cu_seqlens_q.dtype) + row_valid = token_idx < cu_seqlens_q[-1] + pos_in_seq = token_idx - cu_seqlens_q[batch_of_token] + pos_in_seq = torch.where(row_valid, pos_in_seq, torch.zeros_like(pos_in_seq)) + + n_valid_per_row = ((pos_in_seq + 1) // ratio).clamp(max=seq_lens_compressed[batch_of_token]) + n_valid_per_row = torch.where(row_valid, n_valid_per_row, torch.zeros_like(n_valid_per_row)) + offset_per_row = seq_lens_kv[batch_of_token] + + col_idx = ( + torch.arange(max_n_compressed, device=device, dtype=cu_seqlens_q.dtype) + .unsqueeze(0) + .expand(total_q, -1) + ) + valid = col_idx < n_valid_per_row.unsqueeze(1) + matrix = torch.where(valid, col_idx + offset_per_row.unsqueeze(1), torch.full_like(col_idx, -1)) + return matrix.int() + + +def build_cu_seqlens_kv_full( + cu_seqlens_kv: torch.Tensor, cu_seqlens_compressed: torch.Tensor +) -> torch.Tensor: + """Cumulative sequence lengths for the per-segment-concatenated + ``kv_full_thd = cat_per_seg([kv_thd, compressed_kv_thd])``. + + ``kv_full_thd[cu_seqlens_kv_full[b] + i]`` for ``i in [0, seqlen_kv[b])`` + is ``kv_thd[cu_seqlens_kv[b] + i]``; for ``i in [seqlen_kv[b], + seqlen_kv[b] + seqlen_compressed[b])`` it's + ``compressed_kv_thd[cu_seqlens_compressed[b] + (i - seqlen_kv[b])]``. + """ + full_lens = (cu_seqlens_kv[1:] - cu_seqlens_kv[:-1]) + ( + cu_seqlens_compressed[1:] - cu_seqlens_compressed[:-1] + ) + return torch.cat( + [ + torch.zeros(1, dtype=cu_seqlens_kv.dtype, device=cu_seqlens_kv.device), + full_lens.cumsum(0).to(cu_seqlens_kv.dtype), + ] + ) + + +def cat_per_segment( + kv_thd: torch.Tensor, + compressed_kv_thd: Optional[torch.Tensor], + cu_seqlens_kv: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, + cu_seqlens_kv_full: torch.Tensor, +) -> torch.Tensor: + """Build ``kv_full_thd`` by per-segment concatenation of ``kv_thd`` and + ``compressed_kv_thd`` (the THD equivalent of ``torch.cat([kv, + compressed_kv], dim=0)`` in the SBHD path). + + Fully vectorized: computes destination indices for all tokens via + ``batch_of_row`` + offset arithmetic and writes with two indexed + assignments — no Python loop, no GPU→CPU sync. + + Args: + kv_thd: ``(total_kv, *trailing)``. + compressed_kv_thd: ``(total_comp, *trailing)`` or ``None`` if every + segment had ``seqlen < ratio`` (returns ``kv_thd`` unchanged). + cu_seqlens_kv: ``(B+1,)`` int32. + cu_seqlens_compressed:``(B+1,)`` int32. + cu_seqlens_kv_full: ``(B+1,)`` int32 (computed by + :func:`build_cu_seqlens_kv_full`). + + Returns: + ``(total_kv_full, *trailing)`` packed concat. + """ + if compressed_kv_thd is None: + return kv_thd + + total_kv = kv_thd.shape[0] + # NOTE: we deliberately use compressed_kv_thd.shape[0] (capacity, possibly + # padded for CUDA graph capture) rather than cu_seqlens_compressed[-1] (true + # count). The fallback routing on invalid compressed rows (below) writes to + # indices in [total_kv, total_kv_full), so the tail-padding slots *must* + # exist in ``out``. Do not shrink this allocation to true-count without + # also updating the invalid-row routing logic. + total_kv_full = total_kv + compressed_kv_thd.shape[0] + device = kv_thd.device + out_shape = (total_kv_full,) + tuple(kv_thd.shape[1:]) + out = torch.empty(out_shape, dtype=kv_thd.dtype, device=device) + + # KV tokens: dst[i] = cu_full[b] + (i - cu_kv[b]) + batch_of_kv = batch_of_row(cu_seqlens_kv, total_q=total_kv) + src_kv = torch.arange(total_kv, device=device, dtype=cu_seqlens_kv.dtype) + valid_kv = src_kv < cu_seqlens_kv[-1] + dst_kv = cu_seqlens_kv_full[batch_of_kv] + (src_kv - cu_seqlens_kv[batch_of_kv]) + # Invalid (padding) KV rows must be routed to tail-pad slots in + # ``out`` — using ``src_kv`` here is unsafe when ``total_kv > + # cu_seqlens_kv[-1]`` because the padding rows' src indices fall + # inside the valid-kv_full range and race with real-segment writes. + # ``out`` is sized ``total_kv + total_comp_capacity`` so any slot in + # ``[total_kv_full - n_invalid_kv, total_kv_full)`` is reserved + # tail-padding (compressed-invalid uses the same region; duplicate + # writes there are harmless since no valid final index reads them). + dst_kv = torch.where(valid_kv, dst_kv, torch.full_like(dst_kv, total_kv_full - 1)) + out[dst_kv] = kv_thd + + # Compressed tokens: dst[j] = cu_full[b] + kv_len[b] + (j - cu_comp[b]). + # ``compressed_kv_thd`` may be capacity-padded for CUDA graph capture; rows + # beyond ``cu_seqlens_compressed[-1]`` are written to tail padding slots that + # no valid final idx can reference. + total_comp_capacity = compressed_kv_thd.shape[0] + if total_comp_capacity > 0: + kv_lens = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] + src_comp = torch.arange( + total_comp_capacity, device=device, dtype=cu_seqlens_compressed.dtype + ) + batch_of_comp = batch_of_row(cu_seqlens_compressed, total_q=total_comp_capacity) + valid_comp = src_comp < cu_seqlens_compressed[-1] + dst_comp = ( + cu_seqlens_kv_full[batch_of_comp] + + kv_lens[batch_of_comp] + + (src_comp - cu_seqlens_compressed[batch_of_comp]) + ) + dst_comp = torch.where(valid_comp, dst_comp, total_kv + src_comp) + out[dst_comp] = compressed_kv_thd + + return out + + +# --------------------------------------------------------------------------- +# Helper functions for RoPE +# --------------------------------------------------------------------------- + + +def _apply_fused_rope( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + nope_dim: int, + pos_dim: int, + cu_seqlens: Optional[torch.Tensor], + cp_group: torch.distributed.ProcessGroup, +) -> torch.Tensor: + """Apply the fused MLA RoPE kernel with automatic 3-D / 4-D handling.""" + packed_seq = cu_seqlens is not None + + # Strip the dummy batch axis for packed sequences: (total, 1, h, d) → (total, h, d) + squeezed_b = packed_seq and x.dim() == 4 and x.size(1) == 1 + if squeezed_b: + x = x.squeeze(1) + + # Add a dummy head axis for non-packed sequences: (b, s, d) → (b, s, 1, d) + squeeze_head = not packed_seq and x.dim() == 3 + if squeeze_head: + x = x.unsqueeze(-2) + + out = fused_mla_rope_inplace( + x, + cos, + sin, + nope_dim, + pos_dim, + cu_seqlens, + cp_group.rank(), + cp_group.size(), + remove_interleaving=True, + ) + + if squeezed_b: + out = out.unsqueeze(1) + if squeeze_head: + out = out.squeeze(-2) + return out + + +def _apply_unfused_rope( + x: torch.Tensor, + rotary_pos_emb: torch.Tensor, + nope_dim: int, + pos_dim: int, + config: TransformerConfig, + cu_seqlens: Optional[torch.Tensor], + cp_group: torch.distributed.ProcessGroup, + max_seqlen: Optional[int] = None, +) -> torch.Tensor: + """Apply unfused RoPE (split, rotate, concat) with 3-D / 4-D handling. + + DSv4 forces ``mscale=1.0`` — the model relies on Q/KV RMS-norm + + unit-magnitude rotation, not Yarn's concentration factor. + """ + packed_seq = cu_seqlens is not None + + # Drop dummy ``b=1`` from packed 4-D ``(total, 1, h, d)`` callers. + squeezed_b = packed_seq and x.dim() == 4 and x.size(1) == 1 + # Packed 3-D ``(total, 1, d)``: collapse batch and add a temporary head dim. + squeezed_b_3d = packed_seq and x.dim() == 3 and x.size(1) == 1 + if squeezed_b: + x = x.squeeze(1) + elif squeezed_b_3d: + x = x.squeeze(1).unsqueeze(-2) + + # Non-packed 3-D ``(b, s, d)``: add a temporary head dim. + squeeze_head = not packed_seq and x.dim() == 3 + if squeeze_head: + x = x.unsqueeze(-2) + + x_nope, x_pe = torch.split(x, [nope_dim, pos_dim], dim=-1) + x_pe = apply_rotary_pos_emb( + x_pe, + rotary_pos_emb, + config=config, + cu_seqlens=cu_seqlens, + mscale=1.0, + cp_group=cp_group, + mla_rotary_interleaved=True, + mla_output_remove_interleaving=True, + max_seqlen=max_seqlen, + ) + out = torch.cat([x_nope, x_pe], dim=-1) + + if squeezed_b: + out = out.unsqueeze(1) + elif squeezed_b_3d: + out = out.squeeze(-2).unsqueeze(1) + elif squeeze_head: + out = out.squeeze(-2) + return out + + +def _apply_rope( + x: torch.Tensor, + nope_dim: int, + pos_dim: int, + rotary_pos_emb_module: RotaryEmbedding, + config: TransformerConfig, + rotary_seq_len: int, + ratio: int = 1, + cp_group: torch.distributed.ProcessGroup = None, + cu_seqlens: Optional[torch.Tensor] = None, + max_seqlen_rope: Optional[int] = None, +) -> torch.Tensor: + """Apply RoPE to the last ``pos_dim`` dims, leaving the rest unchanged. + + Accepts both 3-D ``[seq, batch, head_dim]`` and 4-D ``[seq, batch, heads, head_dim]`` + inputs. When the input is 3-D a temporary head dimension is inserted for + ``apply_rotary_pos_emb`` and removed before returning. + + Two layouts: + + * **SBHD** (``cu_seqlens=None``): builds a single rotary table of length + ``rotary_seq_len * ratio`` and slices with stride ``ratio``. + * **THD packed** (``cu_seqlens`` supplied): globally strided tables + (``table[:max_total:ratio]``), matching the SBHD approach. + + Args: + max_seqlen_rope: pre-computed ``max(seg_lens) * ratio`` for the + THD + ``ratio > 1`` path (avoids a GPU→CPU sync when the + caller already knows the max original sequence length). + """ + packed_seq = cu_seqlens is not None + + if packed_seq: + if max_seqlen_rope is None: + raise ValueError( + "_apply_rope: max_seqlen_rope is required for THD packed sequences " + "to avoid a GPU→CPU sync that breaks CUDA graph capture." + ) + max_total = max_seqlen_rope + else: + max_total = None + + use_fused = config.apply_rope_fusion + + if use_fused: + # ``mscale=1.0`` keeps the cached cos/sin free of yarn's + # concentration factor so the fused kernel matches the unfused + # split-rotate path (DSv4 "pure rotation" contract). + if packed_seq: + cos, sin = rotary_pos_emb_module.get_cached_cos_sin( + max_total, dtype=x.dtype, packed_seq=True, mscale=1.0 + ) + if ratio > 1: + cos = cos[:max_total:ratio] + sin = sin[:max_total:ratio] + else: + total = rotary_seq_len * ratio if ratio > 1 else rotary_seq_len + cos, sin = rotary_pos_emb_module.get_cached_cos_sin( + total, dtype=x.dtype, packed_seq=False, mscale=1.0 + ) + if ratio > 1: + cos = cos[:total:ratio][:rotary_seq_len] + sin = sin[:total:ratio][:rotary_seq_len] + return _apply_fused_rope(x, cos, sin, nope_dim, pos_dim, cu_seqlens, cp_group) + + # ---- Unfused path: build rotary_pos_emb tensor ---------------------- + if packed_seq: + rope_result = rotary_pos_emb_module(max_total, packed_seq=True) + rotary_pos_emb = rope_result[0] if isinstance(rope_result, tuple) else rope_result + if ratio > 1: + rotary_pos_emb = rotary_pos_emb[:max_total:ratio] + else: + total = rotary_seq_len * ratio if ratio > 1 else rotary_seq_len + rope_result = rotary_pos_emb_module(total, packed_seq=False) + rotary_pos_emb = rope_result[0] if isinstance(rope_result, tuple) else rope_result + if ratio > 1: + rotary_pos_emb = rotary_pos_emb[:total:ratio][:rotary_seq_len] + + # For THD packed sequences ``rotary_pos_emb`` is a single max-length frequency + # table reused per segment, so its (post-stride) length is the max sequence + # length. Passing it as ``max_seqlen`` keeps ``apply_rotary_pos_emb`` on the + # no-global-offset path without a GPU→CPU sync over ``cu_seqlens``. + max_seqlen = rotary_pos_emb.shape[0] if packed_seq else None + return _apply_unfused_rope( + x, rotary_pos_emb, nope_dim, pos_dim, config, cu_seqlens, cp_group, max_seqlen=max_seqlen + ) + + +# --------------------------------------------------------------------------- +# Sparse attention kernel (unfused, differentiable) +# --------------------------------------------------------------------------- + + +def unfused_compressed_sparse_attn( + query: torch.Tensor, + kv_full: torch.Tensor, + attn_sink: torch.Tensor, + topk_indices: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + """Differentiable sparse attention with MQA + learnable attention sink. + Note: the unfused function is mainly for reference, and the performance + and the memory footprint of it is not good for the real scenario. + + Layout is detected from ``query.ndim``: + + * **SBHD** (4-D query): + query ``(sq, b, np, hn)`` multi-head Q. + kv_full ``(n_kv, b, hn)`` single-head MQA KV (original + + compressed concatenated). + topk_indices ``(b, sq, topk)`` int32 **LOCAL per-batch** ids + (``-1`` invalid). + Returns ``(sq, b, np * hn)``. + + * **THD** (3-D query — callers should pre-``squeeze(1)`` the dummy b=1 dim): + query ``(total_q, np, hn)`` packed multi-head Q. + kv_full ``(total_kv, hn)`` packed single-head MQA KV. + topk_indices ``(total_q, topk)`` int32 **flat-global** ids into + ``kv_full`` (``-1`` invalid). + Returns ``(total_q, np * hn)``. + + The math (gather → MQA scores → softmax with sink → weighted sum) is + identical for both layouts; SBHD adds permute / globalize-indices / + unpermute around the call. + + Args: + attn_sink: ``(np,)`` per-head learnable bias for the sink term. + softmax_scale: scalar applied to ``Q · K^T`` before softmax. + """ + is_thd = query.ndim == 3 + + # ----------- Layout-specific input prep ------------------------------- + if is_thd: + q_flat = query # (rows, np, hn) + kv_flat = kv_full # (n_kv, hn) + global_indices = topk_indices # (rows, topk) + else: + sq, b, np_, hn = query.size() + n_kv = kv_full.size(0) + # b-major flatten of query and kv_full. + q_flat = query.permute(1, 0, 2, 3).reshape(b * sq, np_, hn) + kv_flat = kv_full.permute(1, 0, 2).reshape(b * n_kv, hn) + # Globalize topk_indices: ``global = batch_idx * n_kv + local``. + valid = topk_indices >= 0 + batch_ids = torch.arange(b, device=query.device).view(b, 1, 1) + global_indices = torch.where(valid, topk_indices + batch_ids * n_kv, topk_indices).reshape( + b * sq, -1 + ) + + # ----------- Shared core: gather, MQA softmax with sink, sum --------- + rows, np_, hn = q_flat.shape + + safe_indices = global_indices.clamp(min=0).long() + safe_indices_exp = safe_indices.unsqueeze(-1).expand(-1, -1, hn) + kv_gathered = torch.gather( + kv_flat.unsqueeze(0).expand(rows, -1, -1), dim=1, index=safe_indices_exp + ) # (rows, topk, hn) + + q_f = q_flat.float() + kv_g = kv_gathered.float() + scores = torch.einsum("inh,ikh->ink", q_f, kv_g) * softmax_scale # (rows, np, topk) + + invalid_mask = (global_indices < 0).unsqueeze(1) # (rows, 1, topk) + scores = scores.masked_fill(invalid_mask, float("-inf")) + + sink = attn_sink.view(1, np_, 1).float() + scores_max = scores.max(dim=-1, keepdim=True).values + scores_max = torch.max(scores_max, 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("ink,ikh->inh", attn_weights, kv_g) + output = output.to(query.dtype) + + # ----------- Layout-specific output reshape --------------------------- + if is_thd: + return output.reshape(rows, np_ * hn) + 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 +# --------------------------------------------------------------------------- + + +@dataclass +class CompressorSubmodules: + """Submodule specs for CSA and HCA Compressor.""" + + linear_wkv: Union[ModuleSpec, type] = None + linear_wgate: Union[ModuleSpec, type] = None + norm: Union[ModuleSpec, type] = None + + +class Compressor(MegatronModule): + """Gated pooling compressor for CSA and HCA sparse attention. + + Compresses a sequence of tokens into a shorter sequence by pooling groups of + ``compress_ratio`` tokens using learned gated weights. + + For ``compress_ratio == 4``, overlapping compression is used (``coff = 2``). + For ``compress_ratio == 128``, non-overlapping compression is used (``coff = 1``). + + Arbitrary-seqlen handling (same rule for SBHD and THD): + Per-segment ``cutoff = (seqlen // ratio) * ratio = seqlen - (seqlen % ratio)``. + Only the first ``cutoff`` tokens are pooled, producing ``seqlen // ratio`` + compressed entries. The trailing ``seqlen % ratio`` tokens are NOT + compressed and have no compressed-KV representation — they rely on the + sliding window for attention. This matches inference behavior (a + decode token sitting in an incomplete buffer of 1..ratio-1 tokens has + no compressed entry either) and avoids train/inference mismatch from + padding-to-ratio. + + Causal-mask consequence: under the codebase's ``(i+1) // ratio`` + convention, a query token at 0-indexed position ``i`` attends to + ``min((i+1) // ratio, n_compressed_in_segment)`` compressed entries. + The ``clamp`` (in ``get_compress_topk_idxs*`` / kernel-level + ``_indexer_topk_core``) ensures positions in the dropped tail (and + positions in segments shorter than ``ratio``) never index past + ``n_compressed_in_segment``. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: CompressorSubmodules, + compress_ratio: int, + head_dim: int, + rotate: bool = False, + rotary_pos_emb: nn.Module = None, + pg_collection: Optional[ProcessGroupCollection] = None, + name: str | None = None, + ) -> None: + """ + Args: + name (str | None): module instance name passed top-down from its parent module + """ + super().__init__(config=config) + + if pg_collection is None: + # Compatibility fallback for callers not yet passing process groups explicitly. + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + self.pg_collection = pg_collection + + self.compress_ratio = compress_ratio + self.head_dim = head_dim + self.overlap = compress_ratio == 4 + self.coff = 1 + int(self.overlap) + self.rotate = rotate + self.qk_pos_emb_head_dim = config.qk_pos_emb_head_dim + + self.rotary_pos_emb = rotary_pos_emb + + proj_out_dim = self.coff * head_dim + + with get_fp8_disabled_context(config, is_init=True): + self.linear_wkv = build_module( + submodules.linear_wkv, + config.hidden_size, + proj_out_dim, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + name=(name + ".linear_wkv") if name is not None else None, + ) + + self.linear_wgate = build_module( + submodules.linear_wgate, + config.hidden_size, + proj_out_dim, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + name=(name + ".linear_wgate") if name is not None else None, + ) + + # keep to high precision (FP32 in the reference DeepSeek V4 checkpoint) + _ape = torch.empty( + compress_ratio, proj_out_dim, device=torch.cuda.current_device(), dtype=torch.float32 + ) + config.init_method(_ape) + self.ape = mark_keep_in_fp32(nn.Parameter(_ape)) + + norm_config = copy.copy(config) + norm_config.normalization = "RMSNorm" + self.norm = build_module( + submodules.norm, config=norm_config, hidden_size=head_dim, eps=config.layernorm_epsilon + ) + + def _overlap_transform(self, tensor: torch.Tensor, fill_value: float = 0) -> torch.Tensor: + """Apply overlapping window transform for 4x compression. + + Input shape: [n_groups, ratio, b, coff * head_dim] + Output shape: [n_groups, 2 * ratio, b, head_dim] + + Used by the SBHD path where all groups belong to the same sequence. + """ + n_groups, ratio, b_dim, _ = tensor.size() + d = self.head_dim + new_tensor = tensor.new_full((n_groups, 2 * ratio, b_dim, d), fill_value) + new_tensor[:, ratio:] = tensor[:, :, :, d:] + new_tensor[1:, :ratio] = tensor[:-1, :, :, :d] + return new_tensor + + def _overlap_transform_thd( + self, tensor: torch.Tensor, is_first_in_seg: torch.Tensor, fill_value: float = 0 + ) -> torch.Tensor: + """Batched overlapping window transform for THD packed layout. + + Like :meth:`_overlap_transform` but operates on the flat + ``(total_comp, ratio, b, coff * head_dim)`` tensor from all segments + at once. ``is_first_in_seg`` is a ``(total_comp,)`` bool mask that + is ``True`` for each compressed entry that starts a new segment + (i.e. has no predecessor group to pull from). + + Input shape: [total_comp, ratio, b, coff * head_dim] + Output shape: [total_comp, 2 * ratio, b, head_dim] + """ + n, ratio, b_dim, _ = tensor.size() + d = self.head_dim + new_tensor = tensor.new_full((n, 2 * ratio, b_dim, d), fill_value) + new_tensor[:, ratio:] = tensor[:, :, :, d:] + # Previous group's first-half data — shift by 1 along dim-0. + prev_data = torch.roll(tensor[:, :, :, :d], shifts=1, dims=0) + # Zero-fill (or fill_value-fill) segment boundaries. + prev_data[is_first_in_seg] = fill_value + new_tensor[:, :ratio] = prev_data + return new_tensor + + def _project(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Project compressor values and gates outside any enclosing FP8 context.""" + with get_fp8_disabled_context(self.config): + kv, _ = self.linear_wkv(x) + score, _ = self.linear_wgate(x) + return kv, score + + def _forward_sbhd(self, x: torch.Tensor) -> Optional[torch.Tensor]: + """SBHD path. ``x`` is ``(sq, b, hidden_size)``; returns + ``(sq // ratio, b, head_dim)`` or ``None`` when ``sq < ratio``. + """ + sq = x.size(0) + ratio = self.compress_ratio + + if sq < ratio: + return None + + kv, score = self._project(x) # (sq, b, coff * head_dim) + + cutoff = (sq // ratio) * ratio + if cutoff < sq: + kv = kv[:cutoff] + score = score[:cutoff] + n_compressed = cutoff // ratio + + _, b_dim, _ = kv.shape + kv = kv.view(n_compressed, ratio, b_dim, -1) + score = score.view(n_compressed, ratio, b_dim, -1) + score = score + self.ape.view(1, ratio, 1, -1) + if self.overlap: + kv = self._overlap_transform(kv, fill_value=0) + score = self._overlap_transform(score, fill_value=float("-inf")) + weights = torch.softmax(score, dim=1, dtype=torch.float32).to(kv.dtype) + kv = (kv * weights).sum(dim=1) # [n_compressed, b, head_dim] + kv = self.norm(kv.to(x.dtype)) + kv = _apply_rope( + kv, + self.head_dim - self.qk_pos_emb_head_dim, + self.qk_pos_emb_head_dim, + self.rotary_pos_emb, + self.config, + n_compressed, + ratio=ratio, + cp_group=self.pg_collection.cp, + ) + + if self.rotate: + kv = rotate_activation(kv) + return kv + + def _forward_thd( + self, + 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], Optional[torch.Tensor]]: + """THD per-segment compression — fully vectorized. + + Linear projections are token-wise on the flat input. The gated + softmax + reduce is batched across ALL compressed entries from all + segments via a ``(total_compressed, ratio)`` gather index. + + When ``fixed_total_comp`` is supplied the output is padded to that + static capacity so tensor shapes are host-known and CUDA-graph + capturable. Rows beyond the true ``cu_seqlens_compressed[-1]`` + gather from position 0 and are left as tail padding. + + Args: + x: ``(total, 1, hidden_size)`` packed bf16. + cu_seqlens: ``(B+1,)`` int32 cumulative seq lengths + (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. + + Returns: + ``(compressed_thd, cu_seqlens_compressed)`` where + ``compressed_thd`` is ``(total_compressed, 1, head_dim)`` bf16 + (or ``None`` when no sequence has ``seg_len >= ratio`` and + ``fixed_total_comp`` is not set) and + ``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 + + 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 + + # Token-wise projections on the FULL flat input — no boundary issue. + kv, score = self._project(x) # (total, 1, coff * head_dim) + + 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: + 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") + ) + + # Batched softmax + weighted sum — single kernel for all entries. + # (total_comp, [2*]ratio, 1, [coff*]d) → (total_comp, 1, head_dim) + weights = torch.softmax(score_grouped, dim=1, dtype=torch.float32).to(kv_grouped.dtype) + compressed_thd = (kv_grouped * weights).sum(dim=1) + + compressed_thd = self.norm(compressed_thd.to(dtype)) + + 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) + return compressed_thd, cu_seqlens_compressed + + def forward( + self, x: torch.Tensor, packed_seq_params: Optional[PackedSeqParams] = None + ) -> Union[Optional[torch.Tensor], Tuple[Optional[torch.Tensor], torch.Tensor]]: + """Compress hidden states into a shorter KV sequence. + + Two layouts are supported: + + * **SBHD** (default, ``packed_seq_params=None``): ``x`` is + ``(sq, b, hidden_size)``; returns ``(sq // ratio, b, head_dim)`` + (single ``Tensor``) or ``None`` when ``sq < ratio``. + * **THD packed** (``packed_seq_params.qkv_format == 'thd'``): + ``x`` is ``(total, 1, hidden_size)``; returns + ``(compressed_thd, cu_seqlens_compressed)`` (a 2-tuple) so the + caller can build ``kv_full`` per-sequence. ``compressed_thd`` + may be ``None`` when every sequence is shorter than ``ratio``. + """ + nvtx_range_push("compressor") + is_thd = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + if is_thd: + cu_seqlens = ( + packed_seq_params.cu_seqlens_q_padded + if packed_seq_params.cu_seqlens_q_padded is not None + else packed_seq_params.cu_seqlens_q + ) + if packed_seq_params.max_seqlen_q is None: + raise ValueError( + "Compressor: packed_seq_params.max_seqlen_q is required for THD " + "to avoid a GPU→CPU sync that breaks CUDA graph capture." + ) + max_seqlen_q = int(packed_seq_params.max_seqlen_q) + result = self._forward_thd( + x, + cu_seqlens, + max_seqlen_q=max_seqlen_q, + fixed_total_comp=_get_csa_compressed_capacity( + packed_seq_params, self.compress_ratio, x.shape[0] + ), + ) + else: + result = self._forward_sbhd(x) + nvtx_range_pop("compressor") + return result + + +# --------------------------------------------------------------------------- +# CSAIndexer +# --------------------------------------------------------------------------- + + +@dataclass +class CSAIndexerSubmodules: + """Submodule specs for CSAIndexer.""" + + linear_wq_b: Union[ModuleSpec, type] = None + linear_weights_proj: Union[ModuleSpec, type] = None + compressor: Union[ModuleSpec, type] = None + + +class CSAIndexer(MegatronModule): + """Learned top-k retrieval over compressed positions for CSA sparse attention. + + Computes index scores to select the most relevant compressed KV positions for each + query. Reuses the scoring logic from ``DSAIndexer`` (einsum -> relu -> weight -> sum + -> topk) and ``rotate_activation`` (Hadamard transform) from ``dsa.py``. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: CSAIndexerSubmodules, + compress_ratio: int, + rotary_pos_emb: nn.Module = None, + pg_collection: Optional[ProcessGroupCollection] = None, + name: str | None = None, + ) -> None: + """ + Args: + name (str | None): module instance name passed top-down from its parent module + """ + super().__init__(config=config) + + if pg_collection is None: + # Compatibility fallback for callers not yet passing process groups explicitly. + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + self.pg_collection = pg_collection + + self.compress_ratio = compress_ratio + self.hidden_size = config.hidden_size + self.qk_pos_emb_head_dim = config.qk_pos_emb_head_dim + self.q_lora_rank = ( + config.q_lora_rank if config.q_lora_rank is not None else config.hidden_size + ) + + self.index_n_heads = config.dsa_indexer_n_heads + self.index_head_dim = config.dsa_indexer_head_dim + self.index_topk = config.dsa_indexer_topk + + self.softmax_scale: float = self.index_head_dim**-0.5 + + self.rotary_pos_emb = rotary_pos_emb + + # Q projection (FP8 in the reference DeepSeek V4 checkpoint, so it is built + # inside the enclosing fp8_model_init context like the other FP8 weights) + self.linear_wq_b = build_module( + submodules.linear_wq_b, + self.q_lora_rank, + self.index_n_heads * self.index_head_dim, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + name=(name + ".linear_wq_b") if name is not None else None, + ) + + # Weights projection stays in BF16 even under FP8 training (the reference + # DeepSeek V4 checkpoint keeps it in BF16), so build it outside any + # enclosing fp8_model_init context. + with get_fp8_disabled_context(config, is_init=True): + self.linear_weights_proj = build_module( + submodules.linear_weights_proj, + self.hidden_size, + self.index_n_heads, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + name=(name + ".linear_weights_proj") if name is not None else None, + ) + + # Own compressor (smaller head_dim, with Hadamard rotation) + self.compressor = build_module( + submodules.compressor, + config=config, + compress_ratio=compress_ratio, + head_dim=self.index_head_dim, + rotate=True, + rotary_pos_emb=rotary_pos_emb, + pg_collection=pg_collection, + name=(name + ".compressor") if name is not None else None, + ) + + def _project_weights(self, x: torch.Tensor) -> torch.Tensor: + """Project indexer weights outside any enclosing FP8 context.""" + with get_fp8_disabled_context(self.config): + weights, _ = self.linear_weights_proj(x) + return weights + + def forward_before_topk( + self, x: torch.Tensor, qr: torch.Tensor, packed_seq_params: Optional[PackedSeqParams] = None + ) -> Union[ + Tuple[torch.Tensor, torch.Tensor, torch.Tensor], + Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], + ]: + """Compute Q, compressed K, and weights before top-k selection. + + Two layouts: + + * **SBHD** (``packed_seq_params=None``): inputs are ``x (sq, b, h)`` + and ``qr (sq, b, q_lora_rank)``. Returns ``(q, k, weights)``: + ``q (sq, b, n_heads, head_dim)``, ``k (sq // ratio, b, head_dim)``, + ``weights (sq, b, n_heads)``. + * **THD packed** (``packed_seq_params.qkv_format == 'thd'``): + inputs are ``x (total, 1, h)`` and ``qr (total, 1, q_lora_rank)``. + Returns ``(q, k, weights, cu_seqlens_compressed)`` where ``q + (total, 1, n_heads, head_dim)``, ``k (total_comp, 1, head_dim)`` + (``None`` if every sequence is shorter than ``ratio``), + ``weights (total, 1, n_heads)``, and ``cu_seqlens_compressed + (B+1,)`` int32 is the second return value from + ``self.compressor(x, packed_seq_params=...)``. + """ + nvtx_range_push("indexer_before_topk") + + is_thd = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + + sq, bsz, _ = x.size() # in THD: sq = total_q, bsz = 1. + + # ``cu_seqlens_q`` is None for SBHD; ``_apply_rope`` and + # ``self.compressor.forward`` are both layout-aware. + cu_seqlens_q = None + max_seqlen_rope = None + if is_thd: + cu_seqlens_q = ( + packed_seq_params.cu_seqlens_q_padded + if packed_seq_params.cu_seqlens_q_padded is not None + else packed_seq_params.cu_seqlens_q + ) + if packed_seq_params.max_seqlen_q is None: + raise ValueError( + "CSAIndexer: packed_seq_params.max_seqlen_q is required for THD " + "to avoid a GPU→CPU sync that breaks CUDA graph capture." + ) + max_seqlen_rope = int(packed_seq_params.max_seqlen_q) + + # Q path — projection is token-wise so it works for either layout; + # ``_apply_rope`` selects SBHD vs THD packed mode internally + # based on whether ``cu_seqlens`` is supplied. + q, _ = self.linear_wq_b(qr) + q = q.reshape(sq, bsz, self.index_n_heads, self.index_head_dim) + q = _apply_rope( + q, + self.index_head_dim - self.qk_pos_emb_head_dim, + self.qk_pos_emb_head_dim, + self.rotary_pos_emb, + self.config, + rotary_seq_len=sq, + ratio=1, + cp_group=self.pg_collection.cp, + cu_seqlens=cu_seqlens_q, + max_seqlen_rope=max_seqlen_rope, + ) + q = rotate_activation(q) + + # K path: own compressor. SBHD returns ``k``; THD returns the + # 2-tuple ``(k_thd, cu_seqlens_compressed)``. + compressor_out = self.compressor(x, packed_seq_params=packed_seq_params) + + weights = self._project_weights(x) # [sq, b, n_heads] + weights = weights * (self.index_n_heads**-0.5) + + nvtx_range_pop("indexer_before_topk") + if is_thd: + k, cu_seqlens_compressed = compressor_out + return q, k, weights, cu_seqlens_compressed + return q, compressor_out, weights + + def forward( + self, + x: torch.Tensor, + qr: torch.Tensor, + mask: Optional[torch.Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Return (index_scores, topk_indices). + + Two layouts: + + * **SBHD** (default): the original PyTorch reference path using + :func:`fused_qk_topk_naive` with caller-supplied ``mask``. + Returns ``(index_scores (b, sq, sk), topk_indices (b, sq, topk))``. + * **THD packed** (``packed_seq_params.qkv_format == 'thd'``): the + THD analogue that loops per-segment and delegates each one to + :func:`fused_qk_topk_naive` with ``b=1``, then aggregates + per-segment LOCAL top-K ids into a flat + ``(total_q, topk)`` tensor. The per-segment causal + mask is built internally from + :attr:`self.compress_ratio`; ``mask`` is ignored. Returns + ``(None, topk_indices)`` — per-segment scores are not + surfaced because their shapes are heterogeneous and the only + current caller + (:meth:`CompressedSparseAttention._forward_thd` force_unfused + inference) discards them. + """ + nvtx_range_push("indexer") + is_thd = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + if is_thd: + q, k, weights, cu_seqlens_compressed_idx = self.forward_before_topk( + x, qr, packed_seq_params + ) + cu_seqlens_q = ( + packed_seq_params.cu_seqlens_q_padded + if packed_seq_params.cu_seqlens_q_padded is not None + else packed_seq_params.cu_seqlens_q + ) + nvtx_range_push("indexer_qk_topk") + if k is None: + # Every segment is shorter than ``ratio`` → no compressed + # indexer K. Return an all--1 topk so downstream + # consumers treat all positions as invalid. + total_q = q.shape[0] + index_scores = None + topk_indices = torch.full( + (total_q, self.index_topk), -1, dtype=torch.int64, device=q.device + ) + else: + # Squeeze the dummy ``b=1`` dim that ``forward_before_topk`` + # carries (matching the THD shape contract used by the + # cuDNN indexer kernels). + q_thd = q.squeeze(1) + k_thd = k.squeeze(1) + w_thd = weights.squeeze(1) + effective_topk = min(self.index_topk, k_thd.shape[0]) + index_scores, topk_indices = fused_qk_topk_naive_thd( + q_thd, + k_thd, + w_thd, + index_topk=effective_topk, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_compressed_idx, + ratio=self.compress_ratio, + ) + nvtx_range_pop("indexer_qk_topk") + nvtx_range_pop("indexer") + return index_scores, topk_indices + + q, k, weights = self.forward_before_topk(x, qr, packed_seq_params) + nvtx_range_push("indexer_qk_topk") + effective_topk = min(self.index_topk, k.size(0)) + index_scores, topk_indices = fused_qk_topk_naive(q, k, weights, effective_topk, mask) + nvtx_range_pop("indexer_qk_topk") + nvtx_range_pop("indexer") + return index_scores, topk_indices + + +# --------------------------------------------------------------------------- +# CompressedSparseAttention (core attention) +# --------------------------------------------------------------------------- + + +@dataclass +class CompressedSparseAttentionSubmodules: + """Submodule specs for CompressedSparseAttention.""" + + compressor: Union[ModuleSpec, type] = None + indexer: Union[ModuleSpec, type] = None + + +class CompressedSparseAttention(MegatronModule): + """Sparse core attention for CompressedSparseAttention. + + Combines sliding window attention with compressed KV attention. The spec always + provides compressor and indexer submodule specs; this ``__init__`` inspects + ``config.csa_compress_ratios[layer_idx]`` and conditionally builds them: + + * ``ratio == 0``: window-only (compressor and indexer NOT built) + * ``ratio == 4``: window + 4x compressed + learned Indexer (both built) + * ``ratio == 128``: window + 128x compressed, attend to all (compressor built only) + """ + + def __init__( + self, + config: TransformerConfig, + submodules: CompressedSparseAttentionSubmodules, + layer_number: int, + attn_mask_type: AttnMaskType, + attention_type: str, + attention_dropout: Optional[float] = None, + softmax_scale: Optional[float] = None, + k_channels: Optional[int] = None, + v_channels: Optional[int] = None, + cp_comm_type: str = "p2p", + pg_collection: Optional[ProcessGroupCollection] = None, + rotary_pos_emb: nn.Module = None, + compress_ratio: int = 0, + is_mtp_layer: bool = False, + name: str | None = None, + ): + """ + Args: + name (str | None): module instance name passed top-down from its parent module + """ + super().__init__(config=config) + + if pg_collection is None: + # Compatibility fallback for callers not yet passing process groups explicitly. + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + self.pg_collection = pg_collection + + self.layer_number = layer_number + self.config.num_layers if is_mtp_layer else layer_number + self.compress_ratio = compress_ratio + self.window_size = config.csa_window_size + self.v_head_dim = config.v_head_dim + + self.n_local_heads = config.num_attention_heads + + if softmax_scale is None: + softmax_scale = config.v_head_dim**-0.5 + self.softmax_scale = softmax_scale + + self.apply_dsa_kernel_fusion = config.apply_dsa_kernel_fusion + + # Learnable attention sink per head, kept in high precision + # (FP32 in the reference DeepSeek V4 checkpoint) + self.attn_sink = mark_keep_in_fp32( + nn.Parameter(torch.zeros(self.n_local_heads, dtype=torch.float32)) + ) + + # Conditionally build Compressor (ratio > 1) + if self.compress_ratio > 1 and submodules.compressor is not None: + self.compressor = build_module( + submodules.compressor, + config=config, + compress_ratio=self.compress_ratio, + head_dim=config.v_head_dim, + rotate=False, + rotary_pos_emb=rotary_pos_emb, + pg_collection=pg_collection, + name=(name + ".compressor") if name is not None else None, + ) + else: + self.compressor = None + + # Conditionally build Indexer (ratio == 4) + if ( + self.compress_ratio == 4 + and not config.csa_dense_mode + and submodules.indexer is not None + ): + self.indexer = build_module( + submodules.indexer, + config=config, + compress_ratio=self.compress_ratio, + rotary_pos_emb=rotary_pos_emb, + pg_collection=pg_collection, + name=(name + ".indexer") if name is not None else None, + ) + else: + self.indexer = None + + # ------------------------------------------------------------------ + # Private helpers – each owns one logical slice of the forward pass. + # ------------------------------------------------------------------ + + def _build_kv_full( + self, kv: torch.Tensor, x: torch.Tensor + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], int]: + """Concatenate original KV with compressed KV (if applicable). + + Returns: + kv_full: [n_kv, b, v_head_dim] original + compressed KV. + compressed_kv: [n_compressed, b, v_head_dim] or None. + n_compressed: number of compressed positions (0 when unused). + """ + if self.compressor is not None and self.compress_ratio > 1: + compressed_kv = self.compressor(x) + if compressed_kv is not None: + kv_full = torch.cat([kv, compressed_kv], dim=0) + n_compressed = compressed_kv.size(0) + else: + kv_full = kv + compressed_kv = None + n_compressed = 0 + else: + kv_full = kv + compressed_kv = None + n_compressed = 0 + return kv_full, compressed_kv, n_compressed + + def _forward_unfused_csa( + self, + query: torch.Tensor, + x: torch.Tensor, + qr: torch.Tensor, + kv_full: torch.Tensor, + compressed_kv: Optional[torch.Tensor], + n_compressed: int, + offset: int, + window_idxs: torch.Tensor, + packed_seq_params: Optional[PackedSeqParams], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """PyTorch fallback path (no fused kernels). + + Returns ``(output, indexer_loss)``. + """ + sq, b, np, hn = query.size() + indexer_loss = None + + if self.compress_ratio > 1 and n_compressed > 0: + nvtx_range_push("compressed_indices") + if self.indexer is not None: + x_det = x.detach() + qr_det = qr.detach() + + causal_mask = ( + torch.arange(n_compressed, device=x.device).unsqueeze(0).expand(sq, -1) + ) + positions = torch.arange(1, sq + 1, device=x.device).unsqueeze(1) + causal_mask = ( + torch.where(causal_mask >= positions // self.compress_ratio, float("-inf"), 0.0) + .unsqueeze(0) + .expand(b, -1, -1) + ) # [b, sq, n_compressed] + + if self.training and torch.is_grad_enabled(): + q_indexer, k_indexer, weights_indexer = self.indexer.forward_before_topk( + x_det, qr_det, packed_seq_params + ) + indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0) + key_for_loss = compressed_kv.unsqueeze(2).expand(-1, -1, np, -1) + # ``FusedDSAIndexerLoss`` does not accept a separate + # indexer_softmax_scale; apply it here via the + # weights-scaling trick so the effective weights match + # the pre-scale-split behaviour. + weights_for_unfused = weights_indexer.float() * self.indexer.softmax_scale + topk_indices_compressed, indexer_loss = FusedDSAIndexerLoss.apply( + q_indexer, + weights_for_unfused, + k_indexer, + query.detach(), + key_for_loss.detach(), + self.softmax_scale, + min(self.indexer.index_topk, n_compressed), + indexer_loss_coeff, + causal_mask, + getattr(self.config, "dsa_indexer_use_sparse_loss", True), + self.indexer.pg_collection, + # Current main carries packed-sequence metadata before this flag. + None, + None, + None, + None, + self.config.calculate_per_token_loss, + ) + if indexer_loss_coeff > 0: + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=indexer_loss, + layer_number=self.layer_number, + num_layers=self.config.num_layers + (self.config.mtp_num_layers or 0), + ) + else: + _, topk_indices_compressed = self.indexer( + x_det, qr_det, mask=causal_mask, packed_seq_params=packed_seq_params + ) + + n_valid_per_pos = positions // self.compress_ratio # [sq, 1] + valid = (topk_indices_compressed >= 0) & ( + topk_indices_compressed < n_valid_per_pos + ) + compress_topk_idxs = torch.where( + valid, topk_indices_compressed + offset, torch.tensor(-1, device=x.device) + ) + else: + compress_topk_idxs = get_compress_topk_idxs( + self.compress_ratio, b, sq, offset, query.device + ) + + topk_idxs = torch.cat([window_idxs, compress_topk_idxs], dim=-1) + nvtx_range_pop("compressed_indices") + else: + topk_idxs = window_idxs + + topk_idxs = topk_idxs.int() + + nvtx_range_push("sparse_attn_kernel") + output = unfused_compressed_sparse_attn( + query, kv_full, self.attn_sink.float(), topk_idxs, self.softmax_scale + ) + nvtx_range_pop("sparse_attn_kernel") + return output, indexer_loss + + def _forward_fused_no_indexer( + self, + query: torch.Tensor, + kv_full: torch.Tensor, + n_compressed: int, + offset: int, + window_idxs: torch.Tensor, + ) -> torch.Tensor: + """Path A: fused sparse attn with window or deterministic compressed indices.""" + sq, b, np, hn = query.size() + + nvtx_range_push("compressed_indices") + if self.compress_ratio > 1 and n_compressed > 0: + compress_topk_idxs = get_compress_topk_idxs( + self.compress_ratio, b, sq, offset, query.device + ) + flat_idxs, _ = build_flat_topk_idxs(window_idxs, compress_topk_idxs, batch_size=b) + else: + flat_idxs, _ = build_flat_topk_idxs(window_idxs, batch_size=b) + nvtx_range_pop("compressed_indices") + + nvtx_range_push("sparse_attn_kernel") + output = dsa_sparse_attn( + query, kv_full, self.attn_sink.float(), flat_idxs, self.softmax_scale + ) + nvtx_range_pop("sparse_attn_kernel") + return output + + def _forward_fused_indexer_inference( + self, + query: torch.Tensor, + x: torch.Tensor, + qr: torch.Tensor, + kv_full: torch.Tensor, + n_compressed: int, + offset: int, + window_idxs: torch.Tensor, + packed_seq_params: Optional[PackedSeqParams], + ) -> torch.Tensor: + """Path C: separate indexer forward (no loss) + fused sparse attn (compact).""" + b = query.size(1) + + nvtx_range_push("compressed_indices") + x_det = x.detach() + qr_det = qr.detach() + q_indexer, k_indexer, weights_indexer = self.indexer.forward_before_topk( + x_det, qr_det, packed_seq_params + ) + topk_indices_cmp, _ = indexer_topk( + q_indexer, + k_indexer, + weights_indexer, + self.indexer.index_topk, + self.compress_ratio, + indexer_softmax_scale=self.indexer.softmax_scale, + ) + compress_topk_idxs = torch.where(topk_indices_cmp >= 0, topk_indices_cmp + offset, -1) + flat_idxs, flat_tlen = build_flat_topk_idxs( + window_idxs, compress_topk_idxs, batch_size=b, compact=True + ) + nvtx_range_pop("compressed_indices") + + nvtx_range_push("sparse_attn_kernel") + output = dsa_sparse_attn( + query, + kv_full, + self.attn_sink.float(), + flat_idxs, + self.softmax_scale, + topk_length=flat_tlen, + ) + nvtx_range_pop("sparse_attn_kernel") + return output + + def _forward_fused_indexer_training( + self, + query: torch.Tensor, + x: torch.Tensor, + qr: torch.Tensor, + kv_full: torch.Tensor, + n_compressed: int, + offset: int, + window_idxs: torch.Tensor, + packed_seq_params: Optional[PackedSeqParams], + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Path B: fused indexer (with loss) + fused sparse attn. + + Returns ``(output, indexer_loss)``. + """ + nvtx_range_push("compressed_indices") + x_det = x.detach() + qr_det = qr.detach() + q_indexer, k_indexer, weights_indexer = self.indexer.forward_before_topk( + x_det, qr_det, packed_seq_params + ) + nvtx_range_pop("compressed_indices") + + indexer_loss_coeff = self.config.dsa_indexer_loss_coeff or 0.0 + + nvtx_range_push("sparse_attn_kernel") + output, indexer_loss = fused_indexer_sparse_attn( + query, + kv_full, + self.attn_sink.float(), + window_idxs, + q_indexer, + k_indexer, + weights_indexer, + self.indexer.index_topk, + self.compress_ratio, + self.softmax_scale, + self.indexer.softmax_scale, + indexer_loss_coeff, + sparse_loss=getattr(self.config, "dsa_indexer_use_sparse_loss", True), + kv_offset=offset, + calculate_per_token_loss=self.config.calculate_per_token_loss, + ) + nvtx_range_pop("sparse_attn_kernel") + + 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), + ) + return output, indexer_loss + + # ------------------------------------------------------------------ + # Public entry point + # ------------------------------------------------------------------ + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor, + x: torch.Tensor = None, + qr: torch.Tensor = None, + attn_mask_type: AttnMaskType = None, + attention_bias: torch.Tensor = None, + packed_seq_params: PackedSeqParams = None, + boundary_hidden: Optional[torch.Tensor] = None, + boundary_kv: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Forward pass for CompressedSparseAttention. + + Args: + query: [sq, b, np, v_head_dim] + key: [sq, b, 1, v_head_dim] (single-head MQA; head dim squeezed internally) + value: unused (key == value in MQA) + attention_mask: attention mask (may be None for causal). + x: [sq, b, hidden_size] original hidden states. + qr: [sq, b, q_lora_rank] compressed query representation. + + Returns: + output: [sq, b, np * v_head_dim] + """ + nvtx_range_push("compressed_sparse_attn") + + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + 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 + + sq, b, np, hn = query.size() + + kv = key.squeeze(-2) # [sq, b, 1, v_head_dim] -> [sq, b, v_head_dim] + kv_full, compressed_kv, n_compressed = self._build_kv_full(kv, x) + offset = sq # compressed indices start after original positions + window_idxs = get_window_topk_idxs(self.window_size, b, sq, query.device) + + has_indexer_compressed = ( + self.compress_ratio > 1 and n_compressed > 0 and self.indexer is not None + ) + + indexer_loss = None + + if not self.apply_dsa_kernel_fusion: + output, indexer_loss = self._forward_unfused_csa( + query, + x, + qr, + kv_full, + compressed_kv, + n_compressed, + offset, + window_idxs, + packed_seq_params, + ) + elif has_indexer_compressed and self.training and torch.is_grad_enabled(): + output, indexer_loss = self._forward_fused_indexer_training( + query, x, qr, kv_full, n_compressed, offset, window_idxs, packed_seq_params + ) + elif has_indexer_compressed: + output = self._forward_fused_indexer_inference( + query, x, qr, kv_full, n_compressed, offset, window_idxs, packed_seq_params + ) + else: + output = self._forward_fused_no_indexer( + query, kv_full, n_compressed, offset, window_idxs + ) + + if indexer_loss is not None: + output = DSAIndexerLossAutoScaler.apply(output, indexer_loss) + + nvtx_range_pop("compressed_sparse_attn") + return output + + # ------------------------------------------------------------------ + # THD per-path helpers (called from _forward_thd) + # ------------------------------------------------------------------ + + def _forward_unfused_csa_thd( + self, + query: torch.Tensor, + x: torch.Tensor, + qr: torch.Tensor, + kv_full_thd: torch.Tensor, + compressed_kv: Optional[torch.Tensor], + n_compressed_total: int, + np_: int, + total_q: int, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + cu_seqlens_kv_full: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, + window_idxs: torch.Tensor, + max_seqlen_q: int, + max_seqlen_compressed_idx: int, + packed_seq_params: PackedSeqParams, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """PyTorch fallback path for THD (no fused kernels). + + Mirrors :meth:`_forward_unfused_csa` for the SBHD layout. + Returns ``(output, indexer_loss)`` where *output* is + ``(total_q, 1, np * hn)``. + """ + device = query.device + indexer_loss = None + + if self.compress_ratio > 1 and n_compressed_total > 0: + if self.indexer is not None: + x_det = x.detach() + qr_det = qr.detach() + + if self.training and torch.is_grad_enabled(): + q_indexer, k_indexer, weights_indexer, cu_seqlens_compressed_idx = ( + self.indexer.forward_before_topk(x_det, qr_det, packed_seq_params) + ) + if k_indexer is None: + raise RuntimeError( + "CompressedSparseAttention THD unfused Path B requires " + "at least one segment with compressed indexer K." + ) + q_thd = q_indexer.squeeze(1) + w_thd = weights_indexer.squeeze(1) + k_thd = k_indexer.squeeze(1) + + key_for_loss_thd = compressed_kv.unsqueeze(1).expand(-1, np_, -1) + weights_for_unfused = w_thd * self.indexer.softmax_scale + indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0) + + # ``_forward_thd`` (caller) absorbs trailing padded + # tokens into the last ``cu_seqlens_q[-1]`` bucket + # so ``batch_of_row`` doesn't OOB; the Compressor + # does the same to ``cu_seqlens_compressed_idx[-1]``. + # Both are correct for the sparse-attention path, + # but the per-segment indexer-loss loop in + # ``fwd/bwd_fused_indexer_loss_naive_thd`` would + # then iterate a "fake" absorbed-padding segment + # with ``seqlen_k_b < topk`` — triggering a write + # shape mismatch and a downstream ``scatter_`` OOB + # on its ``-1`` entries. + # + # Restore the original (pre-absorption) cu_seqlens + # for the loss path so the segment loop's + # ``if seqlen_k_b == 0: continue`` guard skips the + # padding-only iteration. When no padding exists, + # ``packed_seq_params`` already equals the absorbed + # version and this is a no-op. + # Rebuild compressed cu_seqlens from *unpadded* Q lengths. + # This may disagree with the Indexer's cu_seqlens_compressed_idx + # (which uses padded lengths) for the last segment — that's + # intentional: the extra compressed tokens from padding sit at + # the tail of k_thd and are simply never visited by the loss + # loop, which is correct since they don't represent real data. + # Non-last segments are unaffected (padding absorption only + # extends the final segment). + cu_seqlens_q_for_loss = packed_seq_params.cu_seqlens_q + seg_lens_q = cu_seqlens_q_for_loss[1:] - cu_seqlens_q_for_loss[:-1] + cu_seqlens_compressed_idx_for_loss = torch.cat( + [ + torch.zeros( + 1, + dtype=cu_seqlens_q_for_loss.dtype, + device=cu_seqlens_q_for_loss.device, + ), + (seg_lens_q // self.compress_ratio) + .cumsum(0) + .to(cu_seqlens_q_for_loss.dtype), + ] + ) + topk_indices_cmp, indexer_loss = FusedDSAIndexerLoss.apply( + q_thd, + weights_for_unfused, + k_thd, + query.detach(), + key_for_loss_thd.detach(), + self.softmax_scale, + min(self.indexer.index_topk, max_seqlen_compressed_idx), + indexer_loss_coeff, + None, + getattr(self.config, "dsa_indexer_use_sparse_loss", True), + self.indexer.pg_collection, + None, + None, + None, + None, + self.config.calculate_per_token_loss, + self.config.dsa_indexer_scoring_relu, + cu_seqlens_q_for_loss, + cu_seqlens_compressed_idx_for_loss, + self.compress_ratio, + ) + + if indexer_loss_coeff > 0: + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=indexer_loss, + layer_number=self.layer_number, + num_layers=self.config.num_layers + (self.config.mtp_num_layers or 0), + ) + else: + _, topk_indices_cmp = self.indexer( + x_det, qr_det, mask=None, packed_seq_params=packed_seq_params + ) + + # Shift into per-segment full-KV index space. + if topk_indices_cmp.shape[-1] > 0: + seq_lens_kv = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] + batch_of_token = batch_of_row(cu_seqlens_q, total_q=total_q) + offset_per_row = seq_lens_kv[batch_of_token].unsqueeze(1) + # Per-segment causal post-filter — mirrors the SBHD + # ``_forward_unfused_csa`` post-filter. The training + # indexer (``fwd_fused_indexer_loss_naive_thd``) + # returns RAW per-segment top-K ids without sentinel + # for non-causal picks, so a query at intra-segment + # position ``i`` (0-indexed) may select compressed + # indices ``>= (i+1)//ratio`` whose pre-mask scores + # were ``-inf``; treat those as ``-1`` so the sparse + # attention skips them. + pos_in_seg = ( + torch.arange(total_q, device=device, dtype=cu_seqlens_q.dtype) + - cu_seqlens_q[batch_of_token] + ) + n_valid_per_row = ((pos_in_seg + 1) // self.compress_ratio).unsqueeze(1) + causal_valid = topk_indices_cmp < n_valid_per_row + is_valid = (topk_indices_cmp >= 0) & causal_valid + compress_topk_idxs = torch.where( + is_valid, + topk_indices_cmp + offset_per_row, + torch.full_like(topk_indices_cmp, -1), + ) + else: + compress_topk_idxs = topk_indices_cmp + else: + compress_topk_idxs = get_compress_topk_idxs_thd( + self.compress_ratio, + cu_seqlens_q, + cu_seqlens_kv, + cu_seqlens_compressed, + total_q=total_q, + max_n_compressed=max_seqlen_compressed_idx, + ) + + topk_idxs = torch.cat([window_idxs, compress_topk_idxs], dim=-1) + else: + topk_idxs = window_idxs + + topk_idxs = topk_idxs.int() + + flat_idxs, _ = build_flat_topk_idxs( + topk_idxs, batch_size=-1, cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv_full + ) + + output = unfused_compressed_sparse_attn( + query, kv_full_thd, self.attn_sink.float(), flat_idxs, self.softmax_scale + ) + return output.unsqueeze(1), indexer_loss + + def _forward_fused_no_indexer_thd( + self, + query: torch.Tensor, + kv_full_thd: torch.Tensor, + total_q: int, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + cu_seqlens_kv_full: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, + n_compressed_total: int, + window_idxs: torch.Tensor, + max_seqlen_compressed_idx: int = 0, + ) -> torch.Tensor: + """Path A (THD): fused sparse attn with window or deterministic + compressed indices. + + Returns ``(total_q, 1, np * hn)`` — the attention output. + """ + if self.compress_ratio > 1 and n_compressed_total > 0: + compress_topk_idxs = get_compress_topk_idxs_thd( + self.compress_ratio, + cu_seqlens_q, + cu_seqlens_kv, + cu_seqlens_compressed, + total_q=total_q, + max_n_compressed=max_seqlen_compressed_idx, + ) + flat_idxs, _ = build_flat_topk_idxs( + window_idxs, + compress_topk_idxs, + batch_size=-1, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv_full, + ) + else: + flat_idxs, _ = build_flat_topk_idxs( + window_idxs, + batch_size=-1, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv_full, + ) + + output = dsa_sparse_attn( + query, kv_full_thd, self.attn_sink.float(), flat_idxs, self.softmax_scale, is_thd=True + ) + return output.unsqueeze(1) + + def _forward_fused_indexer_inference_thd( + self, + query: torch.Tensor, + x: torch.Tensor, + qr: torch.Tensor, + kv_full_thd: torch.Tensor, + packed_seq_params: PackedSeqParams, + total_q: int, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + cu_seqlens_kv_full: torch.Tensor, + window_idxs: torch.Tensor, + max_seqlen_q: int, + max_seqlen_compressed_idx: int, + max_seqlen_kv: int, + ) -> torch.Tensor: + """Path C (THD): separate indexer forward (no loss) + fused sparse attn (compact). + + Returns ``(total_q, 1, np * hn)`` — the attention output. + """ + x_det = x.detach() + qr_det = qr.detach() + + q_indexer, k_indexer, weights_indexer, cu_seqlens_compressed_idx = ( + self.indexer.forward_before_topk(x_det, qr_det, packed_seq_params) + ) + q_thd = q_indexer.squeeze(1) + w_thd = weights_indexer.squeeze(1) + if k_indexer is None: + topk_indices_cmp = torch.full((total_q, 0), -1, dtype=torch.int32, device=query.device) + else: + k_thd = k_indexer.squeeze(1) + topk_indices_cmp, _ = indexer_topk( + q_thd, + k_thd, + w_thd, + topk=self.indexer.index_topk, + ratio=self.compress_ratio, + indexer_softmax_scale=self.indexer.softmax_scale, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_compressed_idx, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_compressed_idx, + ) + + # Shift into per-segment full-KV index space. + if topk_indices_cmp.shape[-1] > 0: + seq_lens_kv = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] + batch_of_token = batch_of_row(cu_seqlens_q, total_q=total_q) + offset_per_row = seq_lens_kv[batch_of_token].unsqueeze(1) + compress_topk_idxs = torch.where( + topk_indices_cmp >= 0, + topk_indices_cmp + offset_per_row, + torch.full_like(topk_indices_cmp, -1), + ) + else: + compress_topk_idxs = topk_indices_cmp + + flat_idxs, flat_tlen = build_flat_topk_idxs( + window_idxs, + compress_topk_idxs, + batch_size=-1, + compact=True, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv_full, + ) + output = dsa_sparse_attn( + query, + kv_full_thd, + self.attn_sink.float(), + flat_idxs, + self.softmax_scale, + topk_length=flat_tlen, + is_thd=True, + ) + return output.unsqueeze(1) + + def _forward_fused_indexer_training_thd( + self, + query: torch.Tensor, + x: torch.Tensor, + qr: torch.Tensor, + packed_seq_params: PackedSeqParams, + total_q: int, + np_: int, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + cu_seqlens_kv_full: torch.Tensor, + max_seqlen_q: int, + max_seqlen_compressed_idx: int, + compressed_kv: torch.Tensor, + kv_full_thd: torch.Tensor, + window_idxs: torch.Tensor, + ) -> torch.Tensor: + """Path B (THD): fused indexer (with loss) + fused sparse attn. + + Returns ``(output, indexer_loss)`` where *output* is + ``(total_q, 1, np * hn)``. + """ + sparse_loss = getattr(self.config, "dsa_indexer_use_sparse_loss", True) + indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0) + + x_det = x.detach() + qr_det = qr.detach() + q_indexer, k_indexer, weights_indexer, cu_seqlens_compressed_idx = ( + self.indexer.forward_before_topk(x_det, qr_det, packed_seq_params) + ) + if k_indexer is None: + raise RuntimeError( + "CompressedSparseAttention THD Path B requires at least " + "one segment with compressed indexer K; got none. (Should " + "be unreachable when ``n_compressed_total > 0``.)" + ) + + q_thd = q_indexer.squeeze(1) + w_thd = weights_indexer.squeeze(1) + k_thd = k_indexer.squeeze(1) + + # Supply unpadded cu_seqlens so padding rows are excluded from + # the indexer KL loss (mirrors the unfused path's cu_seqlens_q_for_loss). + # Only pass when they actually differ (by reference or storage) to avoid + # unnecessary mask computation inside the fused kernel. + 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 + + output, indexer_loss = fused_indexer_sparse_attn( + query, + kv_full_thd, + self.attn_sink.float(), + window_idxs, + q_thd, + k_thd, + w_thd, + self.indexer.index_topk, + self.compress_ratio, + self.softmax_scale, + self.indexer.softmax_scale, + indexer_loss_coeff, + sparse_loss=sparse_loss, + kv_offset=0, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + cu_seqlens_kv_full=cu_seqlens_kv_full, + cu_seqlens_compressed_idx=cu_seqlens_compressed_idx, + max_seqlen_q=max_seqlen_q, + max_seqlen_compressed_idx=max_seqlen_compressed_idx, + compressed_kv=compressed_kv, + calculate_per_token_loss=self.config.calculate_per_token_loss, + cu_seqlens_q_unpadded=cu_seqlens_q_unpadded, + ) + + 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), + ) + 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._project_weights(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 + key: torch.Tensor, # (total_kv, 1, 1, hn) packed, MQA + x: torch.Tensor, # (total_q, 1, hidden_size) + qr: torch.Tensor, # (total_q, 1, q_lora_rank) + packed_seq_params: PackedSeqParams, + ) -> torch.Tensor: + """THD-packed branch of :meth:`forward`. See class docstring for layout. + + Performs common setup (shape validation, per-segment compression, + full-KV layout construction, window indices) then dispatches to + one of three per-path helpers: + + * :meth:`_forward_fused_no_indexer_thd` — window-only / window + all-compressed. + * :meth:`_forward_fused_indexer_training_thd` — training + indexer + loss (returns + directly with attached indexer loss). + * :meth:`_forward_fused_indexer_inference_thd` — inference + indexer (no loss). + + Paths A and C return ``compress_topk_idxs`` which are globalized + and fed to the fused/unfused sparse attention in Step 5 below. + """ + # ---- Inputs / shape contract ---------------------------------------- + # query : (total_q, np, hn) multi-head Q (TE THD convention) + # key : (total_kv, 1, 1, hn) packed single-head MQA KV (the + # DSv4 hybrid adds a dummy batch dim to keep the MQA-head + # unsqueeze symmetric with SBHD) + # x, qr : (total_q, 1, *) + total_q, _np, _ = query.shape + + cu_seqlens_q = ( + packed_seq_params.cu_seqlens_q_padded + if packed_seq_params.cu_seqlens_q_padded is not None + else packed_seq_params.cu_seqlens_q + ) + cu_seqlens_kv = ( + packed_seq_params.cu_seqlens_kv_padded + if packed_seq_params.cu_seqlens_kv_padded is not None + else packed_seq_params.cu_seqlens_kv + ) + max_seqlen_q = int(packed_seq_params.max_seqlen_q) + max_seqlen_kv = int(packed_seq_params.max_seqlen_kv) + + # Squeeze the dummy b=1 and MQA head-dim to get the KV-flat layout. + # (key arrives as (total_kv, 1, 1, hn) for MQA.) + kv_thd = key.squeeze(-2).squeeze(1) # (total_kv, hn) + + # ---- Step 2: per-segment compression -------------------------------- + if self.compressor is not None and self.compress_ratio > 1: + compressed_kv, cu_seqlens_compressed = self.compressor( + x, packed_seq_params=packed_seq_params + ) + # compressed_kv is (total_comp, 1, hn) or None + if compressed_kv is not None: + compressed_kv = compressed_kv.squeeze(1) # (total_comp, hn) + n_compressed_total = compressed_kv.shape[0] + else: + n_compressed_total = 0 + else: + compressed_kv = None + cu_seqlens_compressed = torch.zeros_like(cu_seqlens_kv) + n_compressed_total = 0 + + # ---- Build full per-segment-concatenated KV layout ------------------ + cu_seqlens_kv_full = build_cu_seqlens_kv_full(cu_seqlens_kv, cu_seqlens_compressed) + kv_full_thd = cat_per_segment( + kv_thd, compressed_kv, cu_seqlens_kv, cu_seqlens_compressed, cu_seqlens_kv_full + ) + + # ---- Step 3: window indices (per-segment local) --------------------- + window_idxs = get_window_topk_idxs_thd( + self.window_size, cu_seqlens_q, total_q=total_q + ) # (total_q, win_topk) local-to-segment + + # Upper bound on the max compressed-KV length per segment. Not exact + # when segment lengths aren't divisible by compress_ratio, but + # cuDNN/flash kernels tolerate over-estimates (used only for tile sizing). + max_seqlen_compressed_idx = ( + max_seqlen_q // self.compress_ratio if self.compress_ratio > 1 else 0 + ) + + # ---- Step 4: path dispatch -------------------------------------------- + is_training = self.training and torch.is_grad_enabled() + has_indexer = ( + self.compress_ratio > 1 and n_compressed_total > 0 and self.indexer is not None + ) + + indexer_loss = None + + if not self.apply_dsa_kernel_fusion: + output, indexer_loss = self._forward_unfused_csa_thd( + query, + x, + qr, + kv_full_thd, + compressed_kv, + n_compressed_total, + _np, + total_q, + cu_seqlens_q, + cu_seqlens_kv, + cu_seqlens_kv_full, + cu_seqlens_compressed, + window_idxs, + max_seqlen_q, + max_seqlen_compressed_idx, + packed_seq_params, + ) + elif has_indexer and is_training: + output, indexer_loss = self._forward_fused_indexer_training_thd( + query, + x, + qr, + packed_seq_params, + total_q, + _np, + cu_seqlens_q, + cu_seqlens_kv, + cu_seqlens_kv_full, + max_seqlen_q, + max_seqlen_compressed_idx, + compressed_kv, + kv_full_thd, + window_idxs, + ) + elif has_indexer: + output = self._forward_fused_indexer_inference_thd( + query, + x, + qr, + kv_full_thd, + packed_seq_params, + total_q, + cu_seqlens_q, + cu_seqlens_kv, + cu_seqlens_kv_full, + window_idxs, + max_seqlen_q, + max_seqlen_compressed_idx, + max_seqlen_kv, + ) + else: + output = self._forward_fused_no_indexer_thd( + query, + kv_full_thd, + total_q, + cu_seqlens_q, + cu_seqlens_kv, + cu_seqlens_kv_full, + cu_seqlens_compressed, + n_compressed_total, + window_idxs, + max_seqlen_compressed_idx=max_seqlen_compressed_idx, + ) + + if indexer_loss is not None: + output = DSAIndexerLossAutoScaler.apply(output, indexer_loss) + + return output 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..cb7de38b86a --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py @@ -0,0 +1,404 @@ +# 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.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.csa_kernels import indexer_topk + +try: + from megatron.core.fusions.fused_mla_yarn_rope_apply import fused_mla_rope_inplace +except Exception: + fused_mla_rope_inplace = None + +# ============================================================================= +# 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.""" + assert fused_mla_rope_inplace is not None, "Fused MLA RoPE apply is unavailable." + 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/csa_kernels.py b/megatron/core/transformer/experimental_attention_variant/csa_kernels.py new file mode 100644 index 00000000000..76fa75320d7 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/csa_kernels.py @@ -0,0 +1,2074 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +""" +CSA kernel wrappers for Megatron's DSv4 sparse attention. + +Mirrors the three integration paths of the old standalone ``dsa_kernels`` +package, but built on top of + +* :mod:`cudnn.deepseek_sparse_attention` (a.k.a. ``DSA``) — CuTe-DSL backward + + indexer score kernels + TRT-LLM radix top-K, shipped as part of + cuDNN Frontend. +* :mod:`flash_mla` — production sparse-attention forward kernel, expected to + be available as a separate PyPI package. + +Public API (same shape as the old ``dsa_kernels`` package): + +* ``build_flat_topk_idxs`` / ``local_to_global_flat`` — index helpers. +* ``dsa_sparse_attn`` — Path A / Path C step 2, differentiable sparse attention. +* ``indexer_topk`` — Path C inference indexer scoring + top-K. +* ``fused_indexer_sparse_attn`` — Path B training, fused indexer loss + + sparse attention with shared backward. +""" + +from __future__ import annotations + +from functools import lru_cache +from typing import Optional, Tuple + +import torch +from torch import Tensor + +# --------------------------------------------------------------------------- +# Lazy kernel imports +# --------------------------------------------------------------------------- + + +_flash_mla_sparse_fwd = None +_DSA = None + + +def _ensure_flash_mla(): + """Lazily import the FlashMLA sparse-forward kernel. + + FlashMLA ships ``flash_mla_sparse_fwd`` with a multi-head-KV signature; + :func:`_dsa_fwd_flash_mla` below is a thin adapter that unbatches the + DSA-shape inputs and pads ``TopK`` to the alignment expected by + FlashMLA's SM90 / SM100 kernels. + """ + global _flash_mla_sparse_fwd + if _flash_mla_sparse_fwd is not None: + return + + try: + from flash_mla import flash_mla_sparse_fwd as _fwd + except ImportError as e: + raise ImportError( + "FlashMLA is required for DSA sparse attention forward. " + "Install from https://github.com/deepseek-ai/FlashMLA/tree/nv_dev " + "so that `from flash_mla import flash_mla_sparse_fwd` succeeds." + ) from e + _flash_mla_sparse_fwd = _fwd + + +@lru_cache(maxsize=1) +def _get_topk_alignment() -> int: + """Minimum ``TopK`` alignment required by the current GPU architecture. + + * SM90 : dual-warpgroup loop steps by 2 blocks → ``2 * B_TOPK = 128`` + * SM100: single-pipeline loop steps by 1 block → ``B_TOPK`` (64 for + head64, 128 for head128). DSA uses ``D = 512`` which maps to the + head64 kernel path → 64. + """ + sm = torch.cuda.get_device_capability() + if sm[0] >= 10: + return 64 + return 128 + + +def _dsa_fwd_flash_mla( + q: Tensor, + kv: Tensor, + topk_idxs: Tensor, + softmax_scale: float, + d_v: int = 512, + attn_sink: Optional[Tensor] = None, + topk_length: Optional[Tensor] = None, + indexer_topk: int = 0, +) -> Tuple[Tensor, Tensor, Optional[Tensor]]: + """DSA-shaped adapter around :func:`flash_mla.flash_mla_sparse_fwd`. + + Accepts flat (unbatched) tensors with global indices; pads ``TopK`` to + the GPU-specific alignment; returns ``(out, lse, lse_indexer)``. + """ + assert not ( + indexer_topk > 0 and topk_length is not None + ), "indexer_topk > 0 requires non-compact mode (topk_length must be None)" + _ensure_flash_mla() + + _total_S_q, _H, _D = q.shape + TopK = topk_idxs.shape[-1] + topk_align = _get_topk_alignment() + TopK_padded = (TopK + topk_align - 1) // topk_align * topk_align + if TopK_padded != TopK: + pad_width = TopK_padded - TopK + topk_idxs = torch.nn.functional.pad(topk_idxs, (0, pad_width), value=-1) + + kv_3d = kv.unsqueeze(1) # (total_S_kv, 1, D) h_kv=1 + indices = topk_idxs.unsqueeze(1) # (total_S_q, 1, TopK_padded) h_kv=1 + + with torch.cuda.nvtx.range("flash_mla_sparse_fwd"): + res = _flash_mla_sparse_fwd( + q, + kv_3d, + indices, + softmax_scale, + d_v=d_v, + attn_sink=attn_sink, + topk_length=topk_length, + indexer_topk=indexer_topk, + ) + if indexer_topk > 0: + out, _max_logits, lse, lse_indexer = res + else: + out, _max_logits, lse = res + lse_indexer = None + + if indexer_topk > 0: + # When indexer_topk == total TopK, lse_indexer should equal lse but + # the kernel may not snapshot correctly; fall back to lse. + if indexer_topk >= TopK: + return out, lse, lse.clone() + return out, lse, lse_indexer + return out, lse, None + + +def _ensure_dsa_namespace(): + """Lazily import the cudnn-frontend DSA namespace.""" + global _DSA + if _DSA is not None: + return + try: + from cudnn import DSA as _ns + except ImportError as e: + raise ImportError( + "cudnn-frontend DSA namespace not available. Install with " + "`pip install nvidia-cudnn-frontend[cutedsl]`." + ) from e + _DSA = _ns + + +# --------------------------------------------------------------------------- +# Index helpers +# --------------------------------------------------------------------------- + + +def batch_of_row(cu_seqlens_q: Tensor, total_q: Optional[int] = None) -> Tensor: + """For a THD-packed query of length ``total_q``, return a ``(total_q,)`` + int64 tensor where entry ``i`` is the index of the segment that owns + query row ``i`` (i.e. the unique ``b`` with + ``cu_seqlens_q[b] <= i < cu_seqlens_q[b+1]``). + + When ``total_q`` exceeds ``cu_seqlens_q[-1]`` (e.g. after + ``pad_thd_for_cuda_graph`` pads token tensors to a static capacity), + orphan rows are clamped to the last segment so the returned indices + are always in ``[0, B-1]`` and never cause OOB on per-segment arrays. + + Used by every helper that needs to translate between per-row indices + and per-segment cumulative tensors. + + Args: + cu_seqlens_q: ``(B+1,)`` int — cumulative Q lengths. + total_q: optional row count override; defaults to + ``int(cu_seqlens_q[-1].item())`` (forces a GPU→CPU sync). + + Returns: + ``(total_q,)`` int64. + """ + if total_q is None: + total_q = int(cu_seqlens_q[-1].item()) + num_sequences = cu_seqlens_q.shape[0] - 1 + row_idx = torch.arange(total_q, device=cu_seqlens_q.device, dtype=torch.int64) + return torch.bucketize(row_idx, cu_seqlens_q[1:], right=True).clamp( + max=max(num_sequences - 1, 0) + ) + + +def local_to_global_flat( + local_idxs: Tensor, + batch_size: int, + cu_seqlens_q: Optional[Tensor] = None, + cu_seqlens_kv: Optional[Tensor] = None, +) -> Tensor: + """Convert local per-sequence indices to global flat indices. + + Follows the convention used by FlashMLA / SparseAttentionBackward: + flat row order is SBHD ``row[s * B + b]``; global index is + ``local * B + b`` for valid entries and ``-1`` otherwise. + Two layouts are supported: + + * **SBHD-flat (default, ``cu_seqlens_*=None``)** — the convention used + by FlashMLA / SparseAttentionBackward when packing a fixed-shape + batch: flat row order is ``row[s * B + b]``; global index is + ``local * B + b`` for valid entries and ``-1`` otherwise. Inputs + are ``(b, sq, topk)``; outputs are ``(sq*b, topk)``. + * **THD packed (``cu_seqlens_*`` supplied)** — for variable-length + packed sequences. Both ``cu_seqlens_q`` and ``cu_seqlens_kv`` + must be supplied as 1-D int32 tensors of length ``B+1``. Flat row + order is the natural ``(total_q,)`` order; global index is + ``cu_seqlens_kv[batch_of_q] + local`` for valid entries and ``-1`` + otherwise. Inputs are ``(total_q, topk)``; outputs are + ``(total_q, topk)``. + + Args: + local_idxs: SBHD ``(b, sq, topk)`` or THD ``(total_q, topk)`` int. + batch_size: ``B`` (only consulted in the SBHD branch). + cu_seqlens_q: optional 1-D ``(B+1,)`` int32 — when present (with + ``cu_seqlens_kv``), switches to the THD branch. + cu_seqlens_kv: optional 1-D ``(B+1,)`` int32 — same. + + Returns: + ``(sq*b, topk)`` int32 in SBHD mode; ``(total_q, topk)`` int32 in + THD mode. + """ + if (cu_seqlens_q is None) != (cu_seqlens_kv is None): + raise ValueError( + "cu_seqlens_q and cu_seqlens_kv must both be provided for THD, or " + "both None for SBHD." + ) + + if cu_seqlens_q is None: + # ---- SBHD-flat path ------------------------------------------------- + b, sq, topk = local_idxs.shape + assert b == batch_size + + idxs_sb = local_idxs.permute(1, 0, 2).reshape(sq * b, topk) + valid = idxs_sb >= 0 + batch_ids = torch.arange(sq * b, device=local_idxs.device) % b + batch_ids_exp = batch_ids.unsqueeze(1).expand_as(idxs_sb) + idxs_sb = torch.where(valid, idxs_sb * b + batch_ids_exp, idxs_sb) + return idxs_sb.int() + + # ---- THD packed path ---------------------------------------------------- + # Expect ``local_idxs`` to be (total_q, topk). For each row, look up its + # batch index from ``cu_seqlens_q``, then add the corresponding KV offset + # ``cu_seqlens_kv[batch]`` to every valid local index in the row. + if local_idxs.ndim != 2: + raise ValueError(f"THD local_idxs must be 2-D (total_q, topk), got {local_idxs.shape}") + total_q, topk = local_idxs.shape + if cu_seqlens_q.ndim != 1 or cu_seqlens_kv.ndim != 1: + raise ValueError("cu_seqlens_q/kv must be 1-D") + if cu_seqlens_q.shape != cu_seqlens_kv.shape: + raise ValueError( + f"cu_seqlens_q.shape={tuple(cu_seqlens_q.shape)} must equal " + f"cu_seqlens_kv.shape={tuple(cu_seqlens_kv.shape)}" + ) + + row_batch_ids = batch_of_row(cu_seqlens_q, total_q=total_q) + kv_offset = cu_seqlens_kv[row_batch_ids].unsqueeze(1) # (total_q, 1) + valid = local_idxs >= 0 + global_idxs = torch.where(valid, local_idxs + kv_offset, local_idxs) + return global_idxs.int() + + +def build_flat_topk_idxs( + *idx_groups: Tensor, + batch_size: int, + compact: bool = False, + cu_seqlens_q: Optional[Tensor] = None, + cu_seqlens_kv: Optional[Tensor] = None, +) -> Tuple[Tensor, Optional[Tensor]]: + """Combine local per-sequence index groups and convert to flat global form. + + Each *idx_group* contains local per-sequence KV indices (already in + ``kv_full`` index space, i.e. with any compressed-position offset + applied). ``-1`` marks invalid positions. The shape of each group + differs by layout: + + * **SBHD-flat** (``cu_seqlens_*=None``, default): each group is + ``(b, sq, topk_i)``; outputs are ``(sq*b, total_topk)`` (flat + SBHD with row order ``s*B + b``). + * **THD packed** (``cu_seqlens_*`` supplied): each group is + ``(total_q, topk_i)`` with ``total_q = cu_seqlens_q[-1]``; + outputs are ``(total_q, total_topk)``. + + Args: + *idx_groups: one or more index tensors, all of the same layout. + batch_size: ``B`` (only consulted in SBHD). + compact: if True, pack valid entries to the front of each row and + additionally return ``topk_length``; if False, leave as-is and + return ``None``. + cu_seqlens_q: optional 1-D ``(B+1,)`` int32 — selects THD branch. + cu_seqlens_kv: optional 1-D ``(B+1,)`` int32 — selects THD branch. + + Returns: + ``(topk_idxs, topk_length)`` where the first axis of ``topk_idxs`` + is ``sq*b`` (SBHD) or ``total_q`` (THD), and ``topk_length`` is + ``(rows,)`` int32 when ``compact``, else ``None``. + """ + combined = torch.cat(idx_groups, dim=-1) + + # Globalize first, compact second. Both ops are element-wise + + # ``-1``-preserving, so swapping the order is a no-op for correctness; + # globalizing first puts the indices into the same flat row order the + # cuDNN compactify kernel returns its per-row ``length`` in, so no + # extra permute is needed afterward. + global_idxs = local_to_global_flat( + combined, batch_size, cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv + ) + + topk_length_flat = None + if compact: + if global_idxs.is_cuda: + # Fast path: single warp-per-row CuTe DSL kernel from cuDNN's DSA + # namespace. Replaces a stable argsort + gather + sum + permute + # chain with one global-load + global-store per element. + _ensure_dsa_namespace() + res = _DSA.compactify_wrapper(global_idxs) + global_idxs, topk_length_flat = res["indices"], res["topk_length"] + else: + # CPU fallback so the unit tests that exercise this helper without + # CUDA still work. Production callers always go through the CUDA + # path above. + valid_mask = global_idxs >= 0 + sorted_indices = valid_mask.int().argsort(dim=-1, descending=True, stable=True) + global_idxs = global_idxs.gather(-1, sorted_indices) + topk_length_flat = valid_mask.sum(dim=-1).int() + + return global_idxs, topk_length_flat + + +# --------------------------------------------------------------------------- +# Path A + Path C step 2: differentiable sparse attention +# --------------------------------------------------------------------------- + + +class SparseAttnFunc(torch.autograd.Function): + """Sparse attention fwd + bwd on flat tensors. + + Forward uses :mod:`flash_mla`; backward uses cuDNN Frontend's + :attr:`cudnn.DSA.sparse_attention_backward_wrapper`. + """ + + @staticmethod + def forward( + ctx, + q: Tensor, # (total_sq, H, D) bf16 + kv: Tensor, # (total_skv, D) bf16 + attn_sink: Tensor, # (H,) f32 + topk_idxs: Tensor, # (total_sq, TopK) int32 global + topk_length: Optional[Tensor], # (total_sq,) int32 or None + softmax_scale: float, + indexer_topk: int, + ) -> Tuple[Tensor, Tensor, Optional[Tensor]]: + """Run FlashMLA sparse-attention forward and save tensors for backward.""" + out, lse, lse_indexer = _dsa_fwd_flash_mla( + q, + kv, + topk_idxs, + softmax_scale, + attn_sink=attn_sink, + topk_length=topk_length, + indexer_topk=indexer_topk, + ) + + ctx.save_for_backward(q, kv, attn_sink, topk_idxs, out, lse) + ctx.softmax_scale = softmax_scale + ctx.topk_length = topk_length + return out, lse, lse_indexer + + @staticmethod + def backward(ctx, dO, d_lse, d_lse_indexer): + """Compute sparse-attention backward via cuDNN DSA wrapper.""" + _ensure_dsa_namespace() + + q, kv, attn_sink, topk_idxs, out, lse = ctx.saved_tensors + + result = _DSA.sparse_attention_backward_wrapper( + q, + kv, + out, + dO, + lse, + attn_sink, + topk_idxs, + softmax_scale=ctx.softmax_scale, + topk_length=ctx.topk_length, + ) + dq, dkv, d_sink = result["dq"], result["dkv"], result["d_sink"] + return dq, dkv, d_sink, None, None, None, None + + +def dsa_sparse_attn( + query: Tensor, + kv: Tensor, + attn_sink: Tensor, + topk_idxs: Tensor, + softmax_scale: float, + topk_length: Optional[Tensor] = None, + indexer_topk: int = 0, + is_thd: bool = False, +) -> Tensor: + """Sparse attention (Path A / Path C step 2). + + Two layouts: + + * **SBHD** (``is_thd=False``, default): ``query`` is ``(sq, b, np, d)`` + and ``kv`` is ``(skv, b, d)``; the wrapper reshapes them to + ``(sq*b, np, d)`` / ``(skv*b, d)`` before passing to FlashMLA, and + returns ``(sq, b, np * d_v)``. + * **THD packed** (``is_thd=True``): ``query`` is already + ``(total_sq, np, d)`` (3-D) and ``kv`` is ``(total_skv, d)`` (2-D). + No reshape is needed; output is ``(total_sq, np * d_v)`` with a + leading 2-D layout that the caller can fold into its own packed + representation. + + Args: + query: SBHD ``(sq, b, np, d)`` or THD ``(total_sq, np, d)`` bf16. + kv: SBD ``(skv, b, d)`` or THD ``(total_skv, d)`` bf16 (K=V). + attn_sink: ``(np,)`` f32. + topk_idxs: ``(rows, topk)`` int32 — **flat global** indices produced + by :func:`build_flat_topk_idxs` in the matching layout. + softmax_scale: scalar float. + topk_length: ``(rows,)`` int32 — optional compact fast-path. Must be + ``None`` when ``indexer_topk > 0`` (FlashMLA constraint). + indexer_topk: int; ``0`` for Paths A/C, positive for Path B. + is_thd: when True, treat ``query`` and ``kv`` as already-packed + THD tensors and skip the SBHD reshape steps. + + Returns: + SBHD ``(sq, b, np * d_v)`` or THD ``(total_sq, np * d_v)`` bf16. + """ + # Layout-specific input pre-reshape — the kernel always consumes a + # flat ``(rows, np, d)`` query and ``(n_kv, d)`` KV; only the rows + # axis interpretation differs (rows = ``total_sq`` for THD, rows = + # ``sq * b`` for SBHD). ``topk_idxs`` is already a flat ``(rows, k)`` + # tensor in both layouts (built by :func:`build_flat_topk_idxs`). + if is_thd: + if query.ndim != 3: + raise ValueError( + f"THD dsa_sparse_attn expects query of shape " + f"(total_sq, np, d), got {tuple(query.shape)}" + ) + if kv.ndim != 2: + raise ValueError( + f"THD dsa_sparse_attn expects kv of shape (total_skv, d), got {tuple(kv.shape)}" + ) + q_flat, kv_flat = query, kv + else: + sq, b, np_, d = query.shape + skv = kv.shape[0] + q_flat = query.reshape(sq * b, np_, d) + kv_flat = kv.reshape(skv * b, d) + + out_flat, _lse, _lse_indexer = SparseAttnFunc.apply( + q_flat, kv_flat, attn_sink, topk_idxs, topk_length, softmax_scale, indexer_topk + ) # (rows, np, d_v) + + # Layout-specific output reshape: collapse (np, d_v) → (np * d_v), + # then THD stays flat (rows = total_sq); SBHD reflates the (sq, b) axes. + np_, d_v = out_flat.shape[1], out_flat.shape[-1] + if is_thd: + return out_flat.reshape(-1, np_ * d_v) + return out_flat.reshape(sq, b, np_ * d_v) + + +# --------------------------------------------------------------------------- +# Path C inference: indexer scoring + top-K +# --------------------------------------------------------------------------- + + +def _indexer_topk_core( + q: Tensor, + k: Tensor, + w: Tensor, + topk: int, + ratio: int = 4, + *, + cu_seqlens_q: Optional[Tensor] = None, + 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`. + + Wraps cuDNN Frontend's CuTe-DSL indexer-forward kernel. + The pipeline (forward → per-row valid lengths → radix top-K → pad-to-``topk`` → ``topk_length``) + is the same for both layouts; only the input shape glue, valid-length derivation, + and output reshape differ. Selected by ``cu_seqlens_q``. + + BSHD layout (``cu_seqlens_q is None``): + q: ``(b, sq, idx_nh, idx_hd)`` bf16, C-contiguous. + k: ``(b, sk, idx_hd)`` bf16, C-contiguous. + w: ``(b, sq, idx_nh)`` bf16, C-contiguous, **already + ``indexer_softmax_scale``-scaled by the caller**. + Returns: + ``(topk_indices (b, sq, topk) int32, + topk_length (b, sq) int32)`` — invalid slots ``-1``. + + THD packed layout (``cu_seqlens_q is not None``): + q: ``(total_q, idx_nh, idx_hd)`` bf16. + k: ``(total_k, idx_hd)`` bf16. + w: ``(total_q, idx_nh)`` bf16, already scaled. + cu_seqlens_q/kv, max_seqlen_q/kv: standard packed args. + Returns: + ``(topk_indices (total_q, topk) int32, + topk_length (total_q,) int32)`` — per-batch LOCAL ids + in ``[0, seqlen_kv[batch])``; use :func:`local_to_global_flat` + (with ``cu_seqlens_q/kv``) to promote to flat-global ids. + + Two internal entry points besides :func:`indexer_topk`: + + * Path B's ``FusedIndexerSparseAttnFunc.forward`` calls this directly + so the SBHD→BSHD permute can be performed once and reused across + the indexer forward and the score-recompute backward kernels. + """ + is_thd = cu_seqlens_q is not None + device = q.device + + # ---------------- Layout-specific input prep ------------------------ + if is_thd: + if q.ndim != 3: + raise ValueError(f"THD q must be (total_q, idx_nh, idx_hd), got {q.shape}") + if k.ndim != 2: + 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)``. + 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). + scores_flat = scores.contiguous() + sk = int(max_seqlen_kv) + total_q = q.shape[0] + + row_idx = torch.arange(total_q, device=device, dtype=torch.int32) + row_batch_ids = batch_of_row(cu_seqlens_q, total_q=total_q) + row_valid = row_idx < cu_seqlens_q[-1] + pos_in_seq = row_idx - cu_seqlens_q[row_batch_ids] + if q_causal_offsets is not None: + pos_in_seq = pos_in_seq + q_causal_offsets[row_batch_ids] + pos_in_seq = torch.where(row_valid, pos_in_seq, torch.zeros_like(pos_in_seq)) + seqlen_kv_per_row = (cu_seqlens_kv[1:] - cu_seqlens_kv[:-1])[row_batch_ids] + seq_lens = ( + ((pos_in_seq + 1) // ratio).clamp(max=seqlen_kv_per_row).to(torch.int32).contiguous() + ) + seq_lens = torch.where(row_valid, seq_lens, torch.zeros_like(seq_lens)) + else: + if k.shape[1] == 0: + raise ValueError("indexer_topk requires at least one K row.") + + _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" + ] # (b, sq, sk) fp32, -inf on masked positions + b, sq = q.shape[:2] + sk = k.shape[1] + total_q = b * sq + scores_flat = scores.reshape(total_q, sk).contiguous() + + # Per-row valid KV length: ((q_idx + 1) // ratio).clamp(max=sk), + # tiled across the batch axis. + q_idx = torch.arange(sq, device=device) + valid_per_q = ((q_idx + 1) // ratio).clamp(max=sk).to(torch.int32) + seq_lens = valid_per_q.repeat(b) # (b*sq,), row-major over (b, sq) + + # ---------------- Shared: radix top-K + pad-to-topk ----------------- + topk_k = min(topk, sk) + tk_result = _DSA.indexer_top_k_wrapper( + scores_flat, seq_lens, top_k=topk_k, next_n=1, return_val=False + ) + topk_indices = tk_result["indices"] # (total_q, topk_k) int32 + + if topk_k < topk: + pad = torch.full((total_q, topk - topk_k), -1, dtype=torch.int32, device=device) + topk_indices = torch.cat([topk_indices, pad], dim=-1) + + if is_thd: + row_valid = (topk_indices >= 0) & (topk_indices < seq_lens.unsqueeze(1)) + topk_indices = topk_indices.masked_fill(~row_valid, -1) + safe_topk = topk_indices.clamp(min=0, max=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: + return topk_indices.int(), topk_length, scores + return (topk_indices.view(b, sq, topk).int(), topk_length.view(b, sq), scores) + + +def indexer_topk( + q_indexer: Tensor, + k_indexer: Tensor, + weights: Tensor, + topk: int, + ratio: int = 4, + indexer_softmax_scale: float = 1.0, + *, + cu_seqlens_q: Optional[Tensor] = None, + 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). + + Built on cuDNN Frontend's CuTe-DSL indexer forward kernel followed by + TRT-LLM's radix top-K kernel. + + Args: + q_indexer: SBHD ``(sq, b, idx_nh, idx_hd)`` / + THD ``(total_q, idx_nh, idx_hd)`` bf16. + k_indexer: SBHD ``(sk, b, idx_hd)`` / THD ``(total_k, idx_hd)`` bf16. + weights: SBHD ``(sq, b, idx_nh)`` / THD ``(total_q, idx_nh)`` + bf16 — raw (unscaled) weights. + topk: number of top-K indices to select. + ratio: compression ratio for the causal mask. + indexer_softmax_scale: scale applied to the indexer ``Q @ K^T`` + scores (typically ``idx_hd ** -0.5``). Default ``1.0`` means + weights are treated as already-scaled. + cu_seqlens_q: THD only — ``(B+1,)`` int32 CUDA cumulative Q lens. + 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 + — per-batch LOCAL ids into ``k_indexer`` (``-1`` invalid). + THD: ``(topk_indices (total_q, topk), topk_length (total_q,))`` + int32 — per-batch LOCAL ids in ``[0, seqlen_kv[batch])``. + """ + is_thd = cu_seqlens_q is not None + if is_thd and (cu_seqlens_kv is None or max_seqlen_q is None or max_seqlen_kv is None): + raise ValueError( + "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), + # so we push the scale onto the weights tensor (small) instead of the + # score tensor (big). This is uniform across SBHD and THD; in SBHD + # the subsequent permute carries the scaled values into BSHD order. + if indexer_softmax_scale != 1.0: + weights = (weights.float() * indexer_softmax_scale).to(weights.dtype) + + if is_thd: + q, k, w = q_indexer, k_indexer, weights + else: + # SBHD → BSHD permute (one-shot copy each). + q = q_indexer.permute(1, 0, 2, 3).contiguous() + k = k_indexer.permute(1, 0, 2).contiguous() + w = weights.permute(1, 0, 2).contiguous() + + topk_indices, topk_length, _ = _indexer_topk_core( + q, + k, + w, + topk=topk, + ratio=ratio, + cu_seqlens_q=cu_seqlens_q, + 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 + + +# --------------------------------------------------------------------------- +# Path B: fused indexer + sparse attention (training) +# --------------------------------------------------------------------------- + + +_CLIP_PROB_MIN = torch.finfo(torch.float32).tiny # kept compatible w/ cudnn kernel + + +def _thd_to_fake_bshd(*tensors: Tensor) -> Tuple[Tensor, ...]: + """Prepend a B=1 dim to THD tensors for cuDNN wrappers that expect BSHD.""" + return tuple(t.unsqueeze(0) for t in tensors) + + +def _compute_indexer_predict( + q_indexer: Tensor, + k_indexer: Tensor, + weights: Tensor, + topk_indices: Tensor, + qhead_per_kv_head: int, + *, + topk_indices_global: bool = False, +) -> Tensor: + """Compute ``predict`` distribution (softmax over top-K of indexer scores). + + Wraps `cudnn.DSA.sparse_indexer_score_recompute_wrapper`. + This function is not used now, but it is kept for potential future use. + + Two layouts: + + * **BSHD** (default; 4-D q): ``q (B, S_q, H, D)``, ``k (B, S_k, D)``, + ``w (B, S_q, H)``, ``topk (B, S_q, topk)``. + * **THD packed** (3-D q): ``q (total_q, H, D)``, ``k (total_k, D)``, + ``w (total_q, H)``, ``topk (total_q, topk)``. Internally + fake-BSHD'd with ``B=1`` so the wrapper's 4-D-Q shape check + passes; ``topk_indices_global=True`` is required (and enforced) so + the kernel decodes the flat ids directly as positions into the + ``(1*total_k, D)`` view. + + Output shape matches the layout: BSHD ``(B, S_q, topk)`` or + THD ``(total_q, topk)``, fp32 softmax over the top-K axis. + """ + _ensure_dsa_namespace() + is_thd = q_indexer.ndim == 3 + if is_thd: + if not topk_indices_global: + raise ValueError( + "THD ``_compute_indexer_predict`` requires " + "``topk_indices_global=True`` so the kernel addresses K " + "by flat ids over the packed ``(total_k, D)`` buffer." + ) + q_bshd, k_bsd, w_bsh, topk_bst = _thd_to_fake_bshd( + q_indexer, k_indexer, weights, topk_indices + ) + else: + q_bshd, k_bsd, w_bsh, topk_bst = q_indexer, k_indexer, weights, topk_indices + + result = _DSA.sparse_indexer_score_recompute_wrapper( + q_bshd, + k_bsd, + w_bsh, + topk_bst, + qhead_per_kv_head=qhead_per_kv_head, + topk_indices_global=topk_indices_global, + ) + predict = result["predict"] + if is_thd: + predict = predict.squeeze(0) + return predict + + +def _compute_attn_target( + q_attn: Tensor, + k_attn: Tensor, + lse: Tensor, + topk_indices: Tensor, + softmax_scale: float, + qhead_per_kv_head: int, + *, + topk_indices_global: bool = False, +) -> Tensor: + """Compute ``target`` distribution (L1-normalised head-sum softmax). + + Wraps :attr:`cudnn.DSA.sparse_attn_score_recompute_wrapper`. Same + layout convention as :func:`_compute_indexer_predict`: 4-D q is + BSHD; 3-D q is THD and gets fake-BSHD'd with ``B=1`` before the + wrapper call (so the 4-D-Q shape check passes). + """ + _ensure_dsa_namespace() + is_thd = q_attn.ndim == 3 + if is_thd: + if not topk_indices_global: + raise ValueError( + "THD ``_compute_attn_target`` requires " + "``topk_indices_global=True`` so the kernel addresses K " + "by flat ids over the packed ``(total_k, D)`` buffer." + ) + q_bshd, k_bsd, lse_bsh, topk_bst = _thd_to_fake_bshd(q_attn, k_attn, lse, topk_indices) + else: + q_bshd, k_bsd, lse_bsh, topk_bst = q_attn, k_attn, lse, topk_indices + + result = _DSA.sparse_attn_score_recompute_wrapper( + q_bshd, + k_bsd, + lse_bsh, + topk_bst, + softmax_scale, + qhead_per_kv_head=qhead_per_kv_head, + topk_indices_global=topk_indices_global, + ) + target = result["target"] + if is_thd: + target = target.squeeze(0) + return target + + +def _kl_loss_from_target_predict( + target: Tensor, + predict: Tensor, + topk_indices: Tensor, + loss_coeff: float, + calculate_per_token_loss: bool = False, +) -> Tensor: + """KL(target || predict) reduced over ``(B, S_q)`` and scaled by loss_coeff. + + Rows with no valid top-K positions (early query rows with ratio causal + masking) contribute 0 to the loss — the sparse score kernels produce + garbage for those rows, mirroring ``compute_dsa_indexer_loss``'s + ``row_valid`` handling. The default mean is taken over all ``(B, S_q)`` + positions. Per-token-loss mode returns a raw local sum so finalize can + apply the global token divisor. + """ + eps = _CLIP_PROB_MIN + t = target.clamp(min=eps) + p = predict.clamp(min=eps) + kl_per_row = (t * (torch.log(t) - torch.log(p))).sum(dim=-1) # (B, S_q) + + row_valid = (topk_indices >= 0).any(dim=-1) # (B, S_q) + kl_per_row = torch.where(row_valid, kl_per_row, torch.zeros_like(kl_per_row)) + loss = kl_per_row.sum() if calculate_per_token_loss else kl_per_row.mean() + return loss_coeff * loss + + +# --------------------------------------------------------------------------- +# Dense path (``sparse_loss=False``) — full-KV indexer loss +# --------------------------------------------------------------------------- + + +def _compute_dense_indexer_score( + q_indexer: Tensor, + k_indexer: Tensor, + weights: Tensor, + qhead_per_kv_head: int, + indexer_softmax_scale: float, + ratio: int, + *, + cu_seqlens_q: Optional[Tensor] = None, + 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`. + Layout is selected by ``cu_seqlens_*`` kwargs: + + * **BSHD** (``cu_seqlens_*=None``): inputs are 4-D q ``(B, S_q, H, D)``, + 4-D k ``(B, S_k, H_kv, D)``, 3-D w ``(B, S_q, H)``. Outputs are + ``out (B, S_q, S_k)`` + ``denom (B, S_q)``. + * **THD** (``cu_seqlens_*`` supplied): inputs are 3-D q + ``(total_q, H, D)``, 3-D k ``(total_k, H_kv, D)``, 2-D w + ``(total_q, H)``. Outputs are ``out (total_q, max_seqlen_kv)`` + + ``denom (total_q,)``. + + 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, + weights, + qhead_per_kv_head=qhead_per_kv_head, + sm_scale=indexer_softmax_scale, + ratio=ratio, + **kwargs, + ) + return result["out"], result["denom"] + + +_DENSE_ATTN_LSE_CHUNK_MAX_BYTES = 1024 * 1024 * 1024 + + +def _compute_dense_attn_lse_thd( + q_attn_thd: Tensor, + k_attn_thd: Tensor, + softmax_scale: float, + qhead_per_kv_head: int, + ratio: int, + cu_seqlens_q: Tensor, + cu_seqlens_kv: Tensor, + max_seqlen_kv: int, + q_causal_offsets: Optional[Tensor] = None, +) -> Tensor: + """Compute full-KV attention LSE for flattened THD sequences.""" + if q_attn_thd.ndim != 3 or k_attn_thd.ndim != 3: + raise RuntimeError( + "THD attention Q/K must be 3-D, got " + f"q.ndim={q_attn_thd.ndim} and k.ndim={k_attn_thd.ndim}." + ) + if cu_seqlens_q.ndim != 1 or cu_seqlens_kv.ndim != 1: + raise RuntimeError("THD cu_seqlens_q/kv must be 1-D.") + if cu_seqlens_q.shape != cu_seqlens_kv.shape: + raise RuntimeError( + f"THD cu_seqlens shapes must match, got {tuple(cu_seqlens_q.shape)} " + f"and {tuple(cu_seqlens_kv.shape)}." + ) + if q_causal_offsets is not None and ( + q_causal_offsets.ndim != 1 or q_causal_offsets.numel() != cu_seqlens_q.numel() - 1 + ): + raise RuntimeError( + "q_causal_offsets must have one entry per THD segment, got " + f"shape={tuple(q_causal_offsets.shape)} for {cu_seqlens_q.numel() - 1} segments." + ) + + total_q, num_heads, head_dim = q_attn_thd.shape + total_k, num_kv_heads, k_head_dim = k_attn_thd.shape + if head_dim != k_head_dim: + raise RuntimeError( + f"attention Q/K head dimensions must match, got {head_dim} and {k_head_dim}." + ) + if num_heads != num_kv_heads * qhead_per_kv_head: + raise RuntimeError( + "query-head count must equal key-value heads times qhead_per_kv_head, " + f"got h_q={num_heads}, h_kv={num_kv_heads}, " + f"qhead_per_kv_head={qhead_per_kv_head}." + ) + if total_q == 0: + return torch.empty((0, num_heads), device=q_attn_thd.device, dtype=torch.float32) + if max_seqlen_kv <= 0 or total_k == 0: + return torch.full( + (total_q, num_heads), float("-inf"), device=q_attn_thd.device, dtype=torch.float32 + ) + + row_batch_ids = batch_of_row(cu_seqlens_q, total_q=total_q) + row_positions = torch.arange(total_q, device=q_attn_thd.device, dtype=torch.int64) + row_positions = row_positions - cu_seqlens_q[row_batch_ids].to(torch.int64) + if q_causal_offsets is not None: + row_positions = row_positions + q_causal_offsets[row_batch_ids].to(torch.int64) + kv_starts = cu_seqlens_kv[row_batch_ids].to(torch.int64) + kv_lengths = (cu_seqlens_kv[1:] - cu_seqlens_kv[:-1])[row_batch_ids].to(torch.int64) + valid_k_per_row = torch.minimum((row_positions + 1) // ratio, kv_lengths) + key_positions = torch.arange(max_seqlen_kv, device=q_attn_thd.device, dtype=torch.int64) + + lse = torch.empty((total_q, num_heads), device=q_attn_thd.device, dtype=torch.float32) + # Advanced indexing materializes one K vector per query row. Bound both + # that buffer and the score tensor, rather than accounting for scores only. + bytes_per_row = max_seqlen_kv * (head_dim + qhead_per_kv_head) * 4 + chunk_rows = min(total_q, max(1, _DENSE_ATTN_LSE_CHUNK_MAX_BYTES // max(1, bytes_per_row))) + + for kv_head in range(num_kv_heads): + head_start = kv_head * qhead_per_kv_head + head_end = head_start + qhead_per_kv_head + for q_start in range(0, total_q, chunk_rows): + q_end = min(q_start + chunk_rows, total_q) + global_k = kv_starts[q_start:q_end].unsqueeze(1) + key_positions.unsqueeze(0) + global_k = global_k.clamp(max=total_k - 1) + k_group = k_attn_thd[global_k, kv_head, :].float() + q_group = q_attn_thd[q_start:q_end, head_start:head_end, :].float() + scores = torch.einsum("qhd,qkd->qhk", q_group, k_group) * softmax_scale + valid = key_positions.unsqueeze(0) < valid_k_per_row[q_start:q_end].unsqueeze(1) + scores.masked_fill_(~valid.unsqueeze(1), float("-inf")) + lse[q_start:q_end, head_start:head_end] = torch.logsumexp(scores, dim=-1) + + return lse.contiguous() + + +def _compute_dense_attn_lse( + q_attn_bshd: Tensor, + k_attn_bshd: Tensor, + softmax_scale: float, + qhead_per_kv_head: int, + ratio: int, + *, + cu_seqlens_q: Optional[Tensor] = None, + cu_seqlens_kv: Optional[Tensor] = None, + max_seqlen_kv: Optional[int] = None, + q_causal_offsets: Optional[Tensor] = None, +) -> Tensor: + """Compute full-KV per-query-head attention LSE for BSHD or THD.""" + if (cu_seqlens_q is None) != (cu_seqlens_kv is None): + raise ValueError( + "cu_seqlens_q and cu_seqlens_kv must both be provided for THD, or " + "both None for BSHD." + ) + if cu_seqlens_q is not None: + if max_seqlen_kv is None: + raise ValueError("max_seqlen_kv is required for THD dense attention LSE.") + return _compute_dense_attn_lse_thd( + q_attn_bshd, + k_attn_bshd, + softmax_scale, + qhead_per_kv_head, + ratio, + cu_seqlens_q, + cu_seqlens_kv, + int(max_seqlen_kv), + q_causal_offsets, + ) + + if q_causal_offsets is not None: + raise ValueError("q_causal_offsets is only supported in THD mode.") + + b, sq, num_heads, _ = q_attn_bshd.shape + kb, sk, num_kv_heads, _ = k_attn_bshd.shape + if kb != b: + raise RuntimeError(f"attention Q/K batch sizes must match, got {b} and {kb}.") + if num_heads != num_kv_heads * qhead_per_kv_head: + raise RuntimeError( + "query-head count must equal key-value heads times qhead_per_kv_head, " + f"got h_q={num_heads}, h_kv={num_kv_heads}, " + f"qhead_per_kv_head={qhead_per_kv_head}." + ) + + lse = torch.empty((b, sq, num_heads), device=q_attn_bshd.device, dtype=torch.float32) + key_positions = torch.arange(sk, device=q_attn_bshd.device).view(1, sk) + q_positions = torch.arange(sq, device=q_attn_bshd.device) + seq_lens = ((q_positions + 1) // ratio).clamp(max=sk) + score_bytes_per_row = b * qhead_per_kv_head * sk * torch.finfo(torch.float32).bits // 8 + chunk_rows = min(sq, max(1, _DENSE_ATTN_LSE_CHUNK_MAX_BYTES // max(1, score_bytes_per_row))) + + for kv_head in range(num_kv_heads): + head_start = kv_head * qhead_per_kv_head + head_end = head_start + qhead_per_kv_head + k_group = k_attn_bshd[:, :, kv_head, :].float() + for q_start in range(0, sq, chunk_rows): + q_end = min(q_start + chunk_rows, sq) + q_group = q_attn_bshd[:, q_start:q_end, head_start:head_end, :].float() + scores = torch.einsum("bqhd,bkd->bqhk", q_group, k_group) * softmax_scale + valid = key_positions < seq_lens[q_start:q_end].view(-1, 1) + scores.masked_fill_(~valid.view(1, q_end - q_start, 1, sk), float("-inf")) + lse[:, q_start:q_end, head_start:head_end] = torch.logsumexp(scores, dim=-1) + + return lse.contiguous() + + +def _compute_dense_attn_score( + q_attn: Tensor, + k_attn: Tensor, + lse: Tensor, + qhead_per_kv_head: int, + softmax_scale: float, + ratio: int, + *, + cu_seqlens_q: Optional[Tensor] = None, + 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). + + Wraps :attr:`cudnn.DSA.dense_attn_score_recompute_wrapper`. Same + 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, + lse, + softmax_scale, + qhead_per_kv_head=qhead_per_kv_head, + ratio=ratio, + **kwargs, + ) + return result["out"], result["denom"] + + +def _kl_loss_from_dense_scores( + attn_score: Tensor, + attn_l1norm: Tensor, + index_score: Tensor, + index_lse: Tensor, + loss_coeff: float, + calculate_per_token_loss: bool = False, +) -> Tensor: + """KL(target || predict) over the **full** KV axis, averaged over rows. + + Derives ``target = attn_score / attn_l1norm`` (L1-normalised, matches + ``compute_dsa_indexer_loss``'s ``attention_scores / sum`` step) and + ``log_predict = index_score - index_lse`` (LSE-normalised log-softmax), + then computes ``KL = sum_k target * (log target - log predict)`` and + scales by ``loss_coeff``. + + Layout-agnostic: works for both BSHD inputs (shapes + ``attn_score (B, S_q, S_k)``, ``attn_l1norm (B, S_q)``, …) and THD + inputs (shapes ``attn_score (total_q, max_seqlen_kv)``, + ``attn_l1norm (total_q,)``, …). The final ``.mean()`` averages over + all rows in either case. + + Rows where the kernel's ``ratio`` causal mask leaves no valid KV + position have ``attn_l1norm <= 0`` (L1) or ``index_lse == -inf`` + (LSE); those rows contribute 0 to the loss — the same ``row_valid`` + semantics as the reference ``compute_dsa_indexer_loss``. + """ + eps = _CLIP_PROB_MIN + # row_valid: rows with at least one un-masked KV position. + row_valid = (attn_l1norm > eps) & torch.isfinite(index_lse) + + # Safe denoms: replace invalid rows with a finite value so target / + # log-predict don't produce NaN; the row mask zeroes their KL below. + safe_l1 = attn_l1norm.clamp(min=eps) + safe_lse = torch.where(row_valid, index_lse, torch.zeros_like(index_lse)) + + target = attn_score / safe_l1.unsqueeze(-1) + target_clamped = target.clamp(min=eps) + # Per-position validity: the indexer-score kernel emits -inf at + # ratio-masked positions; those contribute 0 to KL by the + # ``0 * log(0/p) = 0`` convention. Without this gate, the eps-clamp + # on target makes the term ``eps * (log eps - (-inf)) = +inf``. + position_valid = torch.isfinite(index_score) + safe_index_score = torch.where(position_valid, index_score, torch.zeros_like(index_score)) + log_predict = safe_index_score - safe_lse.unsqueeze(-1) + + kl_terms = target_clamped * (torch.log(target_clamped) - log_predict) + kl_terms = torch.where(position_valid, kl_terms, torch.zeros_like(kl_terms)) + kl_per_row = kl_terms.sum(dim=-1) # (B, S_q) + kl_per_row = torch.where(row_valid, kl_per_row, torch.zeros_like(kl_per_row)) + loss = kl_per_row.sum() if calculate_per_token_loss else kl_per_row.mean() + return loss_coeff * loss + + +class FusedIndexerSparseAttnFunc(torch.autograd.Function): + """Path B: fused indexer (+KL loss) + sparse attention in one autograd. + + Differentiable w.r.t. ``query``, ``kv_full``, ``attn_sink``, + ``q_indexer``, ``k_indexer``, ``weights``. + + Layout is selected by the ``cu_seqlens_q`` kwarg passed to + :func:`fused_indexer_sparse_attn`: + + * **SBHD** (``cu_seqlens_q is None``): inputs carry an explicit batch + axis; the indexer pipeline runs in BSHD (after one SBHD→BSHD + permute). + * **THD packed** (``cu_seqlens_q`` supplied): inputs are flat + packed-sequence tensors; the indexer pipeline runs directly on + ``(total_q, …)`` / ``(total_kv, …)`` shapes with + ``cu_seqlens_q/kv`` forwarded to every layout-aware kernel + (``_indexer_topk_core`` THD branch, ``local_to_global_flat`` THD + branch, ``_compute_dense_*_score`` THD branch, + ``dense_indexer_backward_wrapper`` with ``cu_seqlens_q/k``). + + Two indexer-loss variants, selected by the ``sparse_loss`` argument + (matches ``compute_dsa_indexer_loss`` in the reference ``dsa.py``): + + * **Sparse loss** (``sparse_loss=True``) — KL is computed only over + the top-K KV positions the indexer has selected. + **Supports both SBHD and THD.** + * **Dense loss** (``sparse_loss=False``, the default) — KL is + computed over *all* causally valid KV positions. + **Supports both SBHD and THD.** + + The indexer backward is eagerly computed in the forward pass with + ``grad_loss=1.0``; the actual backward simply scales the + pre-computed gradients by ``grad_loss``. + + Both variants share the FlashMLA sparse-attention forward + the + cuDNN sparse-attn backward (both of which are layout-agnostic — the + flat shape they require is what the THD branch already passes + directly, and what the SBHD branch reshapes into). + """ + + @staticmethod + def forward( + ctx, + # Sparse attn inputs (differentiable) + query: Tensor, # SBHD (sq, b, np, d) / THD (total_q, np, d) + kv_full: Tensor, # SBHD (skv, b, d) / THD (total_kv_full, d) + attn_sink: Tensor, # (np,) f32 + # Window indices (not differentiable) + window_idxs: Tensor, # SBHD (b, sq, win_topk) / THD (total_q, win_topk) + # Indexer inputs (differentiable) + q_indexer: Tensor, # SBHD (sq, b, idx_nh, idx_hd) / THD (total_q, idx_nh, idx_hd) + k_indexer: Tensor, # SBHD (n_comp, b, idx_hd) / THD (total_comp_idx, idx_hd) + weights: Tensor, # SBHD (sq, b, idx_nh) / THD (total_q, idx_nh) — raw (unscaled) + # Scalars + indexer_topk: int, + ratio: int, + softmax_scale: float, + indexer_softmax_scale: float, + loss_coeff: float, + sparse_loss: bool, + kv_offset: int, # SBHD only — start of compressed region in kv_full + calculate_per_token_loss: bool, + # THD packed-sequence args (all None for SBHD; all required for THD) + cu_seqlens_q: Optional[Tensor], + cu_seqlens_kv: Optional[Tensor], # original (uncompressed) KV cu_seqlens + cu_seqlens_kv_full: Optional[Tensor], # original + compressed concat'd cu_seqlens + cu_seqlens_compressed_idx: Optional[Tensor], # indexer K cu_seqlens (== compressor's) + max_seqlen_q: Optional[int], + max_seqlen_compressed_idx: Optional[int], # indexer K max + compressed_kv: Optional[Tensor] = None, # THD only — pre-packed compressed KV + cu_seqlens_q_unpadded: Optional[Tensor] = None, # THD only — unpadded Q cu_seqlens + ) -> Tuple[Tensor, Tensor]: + """Fused forward: indexer scoring, sparse attention, KL loss, and indexer backward.""" + _ensure_dsa_namespace() + + is_thd = cu_seqlens_q is not None + + # ---- Layout-specific input prep -------------------------------------- + # SBHD: permute SBHD→BSHD once and reuse the BSHD tensors for indexer + # forward, dense score helpers, and the indexer backward. + # THD: skip the permute; tensors are already flat. + if is_thd: + total_q = q_indexer.shape[0] + idx_nh = q_indexer.shape[1] + np_, d = query.shape[1], query.shape[2] + + q_indexer_flat = q_indexer + k_indexer_flat = k_indexer + w_indexer = weights + else: + sq, b, np_, d = query.shape + skv = kv_full.shape[0] + idx_nh = q_indexer.shape[2] + + q_indexer_flat = q_indexer.permute(1, 0, 2, 3).contiguous() + k_indexer_flat = k_indexer.permute(1, 0, 2).contiguous() + w_indexer = weights.permute(1, 0, 2).contiguous() + + # ``indexer_softmax_scale`` is applied via the + # ``relu(c·x) = c·relu(x)`` trick (the cudnn kernel does the relu), + # so we push the scale onto the weights tensor (small) instead of the + # score tensor (big). This is uniform across SBHD and THD; in SBHD + # the subsequent permute carries the scaled values into BSHD order. + if indexer_softmax_scale != 1.0: + w_indexer_scaled = (w_indexer.float() * indexer_softmax_scale).to(w_indexer.dtype) + else: + w_indexer_scaled = w_indexer + + # ---- 2. Indexer scoring + top-K (with scores retained). --------------- + # Pass the original ``indexer_topk`` (not min(indexer_topk, n_comp)) so + # that the output is always padded to a fixed size. flash_mla_sparse_fwd + # requires a consistent TopK dimension; _indexer_topk_core handles the + # case where sk < topk internally (selects min(topk, sk) values, then + # pads to topk with -1). + topk_indices_cmp, _, indexer_scores = _indexer_topk_core( + q_indexer_flat, + k_indexer_flat, + w_indexer_scaled, + indexer_topk, + ratio, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_compressed_idx, + max_seqlen_q=int(max_seqlen_q) if max_seqlen_q is not None else None, + max_seqlen_kv=( + int(max_seqlen_compressed_idx) if max_seqlen_compressed_idx is not None else None + ), + ) + + # ---- 3. Combine indices (indexer first, then window) + globalize. ---- + if is_thd: + row_batch_ids = batch_of_row(cu_seqlens_q, total_q=total_q) + offset_per_row = ( + (cu_seqlens_kv[1:] - cu_seqlens_kv[:-1])[row_batch_ids].unsqueeze(1).to(torch.int32) + ) + compress_topk_idxs = torch.where( + topk_indices_cmp >= 0, + topk_indices_cmp + offset_per_row, + torch.full_like(topk_indices_cmp, -1), + ) + combined_local = torch.cat( + [compress_topk_idxs, window_idxs], dim=-1 + ) # (total_q, indexer_topk + win_topk) + global_idxs = local_to_global_flat( + combined_local, + batch_size=-1, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv_full, + ) + else: + compress_topk_idxs = torch.where( + topk_indices_cmp >= 0, topk_indices_cmp + kv_offset, -1 + ) + combined_local = torch.cat([compress_topk_idxs, window_idxs], dim=-1) + global_idxs = local_to_global_flat(combined_local, b) + + # ---- 4. FlashMLA forward (flat layout for both SBHD and THD). -------- + if is_thd: + q_flat = query + kv_flat = kv_full + else: + q_flat = query.reshape(sq * b, np_, d) + kv_flat = kv_full.reshape(skv * b, d) + out_flat, lse, lse_indexer = _dsa_fwd_flash_mla( + q_flat, + kv_flat, + global_idxs, + softmax_scale, + attn_sink=attn_sink, + topk_length=None, + indexer_topk=indexer_topk, + ) + + # ---- 4b. Derive padding-row mask for loss exclusion. ----------------- + # When CUDA-graph padding makes cu_seqlens_q cover all total_q rows + # (including padding), cu_seqlens_q_unpadded supplies the true + # boundaries. Padding rows must not contribute to the indexer KL + # loss or backward gradients — only the sparse-attention output + # needs them for static-shape compatibility. + # The caller only passes cu_seqlens_q_unpadded when it differs from + # cu_seqlens_q (checked via data_ptr), so no GPU→CPU sync is needed. + padding_row_mask: Optional[Tensor] = None # True = padding (excluded from loss) + if is_thd and cu_seqlens_q_unpadded is not None: + real_seg_lens = cu_seqlens_q_unpadded[1:] - cu_seqlens_q_unpadded[:-1] + row_idx = torch.arange(total_q, device=query.device, dtype=torch.int32) + row_batch_ids = batch_of_row(cu_seqlens_q, total_q=total_q) + pos_in_seg = row_idx - cu_seqlens_q[row_batch_ids].to(torch.int32) + # Rows whose intra-segment position >= real segment length + # are padding (including rows in a dummy trailing segment + # whose real length is 0). + real_len_per_row = real_seg_lens[row_batch_ids].to(torch.int32) + padding_row_mask = pos_in_seg >= real_len_per_row + + # ---- 5. Derive predict from indexer_scores, compute target. ---------- + # Layout-specific attn tensors (detached — loss is not differentiable + # through them). + if is_thd: + assert compressed_kv is not None, "compressed_kv is required for THD" + q_attn_det = query.detach() + k_attn_compressed_det = compressed_kv.detach() + lse_indexer_det = lse_indexer.detach() + else: + q_attn_det = query.detach().permute(1, 0, 2, 3).contiguous() + k_attn_compressed_det = kv_full[kv_offset:].detach().permute(1, 0, 2).contiguous() + lse_indexer_det = lse_indexer.reshape(sq, b, np_).permute(1, 0, 2) + + # Invalidate padding rows for the loss/backward path. The sparse + # attention (steps 3-4) has already built global_idxs from the + # original topk_indices_cmp, so this mutation only affects steps 5-7. + if padding_row_mask is not None: + topk_indices_cmp = topk_indices_cmp.clone() + topk_indices_cmp[padding_row_mask] = -1 + indexer_scores = indexer_scores.clone() + indexer_scores[padding_row_mask] = float('-inf') + + if sparse_loss: + # Derive predict: gather topk scores from indexer_scores → softmax. + safe_indices = topk_indices_cmp.clamp(min=0).long() + gathered_scores = torch.gather(indexer_scores, dim=-1, index=safe_indices) + gathered_scores = torch.where( + topk_indices_cmp >= 0, gathered_scores, torch.finfo(torch.float32).min + ) + predict = torch.softmax(gathered_scores, dim=-1) + + # THD: _compute_attn_target's kernel addresses K by flat ids over + # the packed (total_k, D) buffer, so promote per-segment-local + # indices to flat-global against cu_seqlens_compressed_idx. + if is_thd: + topk_for_target = local_to_global_flat( + topk_indices_cmp, + batch_size=-1, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_compressed_idx, + ) + else: + topk_for_target = topk_indices_cmp + + target = _compute_attn_target( + q_attn_det, + k_attn_compressed_det, + lse_indexer_det, + topk_for_target, + softmax_scale, + qhead_per_kv_head=np_, + topk_indices_global=is_thd, + ) + + if loss_coeff > 0: + indexer_loss = _kl_loss_from_target_predict( + target, predict, topk_indices_cmp, loss_coeff, calculate_per_token_loss + ) + else: + indexer_loss = torch.zeros((), device=query.device, dtype=torch.float32) + else: + index_score = indexer_scores + index_lse = torch.logsumexp(indexer_scores, dim=-1) + + k_unsqueeze_dim = 1 if is_thd else 2 + dense_attn_kwargs = {} + if is_thd: + dense_attn_kwargs = dict( + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_compressed_idx, + max_seqlen_q=int(max_seqlen_q), + max_seqlen_kv=int(max_seqlen_compressed_idx), + ) + dense_lse = _compute_dense_attn_lse( + q_attn_det, + k_attn_compressed_det.unsqueeze(k_unsqueeze_dim), + softmax_scale, + qhead_per_kv_head=np_, + ratio=ratio, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_compressed_idx, + max_seqlen_kv=max_seqlen_compressed_idx, + ) + attn_score, attn_l1norm = _compute_dense_attn_score( + q_attn_det, + k_attn_compressed_det.unsqueeze(k_unsqueeze_dim), + dense_lse, + qhead_per_kv_head=np_, + softmax_scale=softmax_scale, + 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( + attn_score, + attn_l1norm, + index_score, + index_lse, + loss_coeff, + calculate_per_token_loss, + ) + else: + indexer_loss = torch.zeros((), device=query.device, dtype=torch.float32) + + # ---- 6. Eagerly compute indexer backward (grad_loss=1). ------------ + # The actual grad_loss scaling is deferred to backward (when + # DSAIndexerLossAutoScaler provides the correct scale). + # Use total_q (not real token count) for the loss coefficient even + # when padding rows are masked. The cuDNN kernel divides by total_q + # internally; since masked rows contribute 0, multiplying back by + # total_q still yields the correct real-token sum — and avoids a + # GPU→CPU sync that would break CUDA graph capture. + indexer_loss_coeff = loss_coeff + if calculate_per_token_loss: + indexer_loss_coeff = loss_coeff * (total_q if is_thd else b * sq) + + unit_grad_loss = torch.ones((), device=query.device, dtype=torch.float32) + + if loss_coeff > 0: + if sparse_loss: + attn_score_for_bwd = target.clone() + index_score_for_bwd = predict.clone() + if is_thd: + topk_indices_cmp_global = local_to_global_flat( + topk_indices_cmp, + batch_size=-1, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_compressed_idx, + ) + bwd_q, bwd_w, bwd_k, bwd_attn, bwd_idx, bwd_topk = _thd_to_fake_bshd( + q_indexer_flat, + w_indexer, + k_indexer_flat, + attn_score_for_bwd, + index_score_for_bwd, + topk_indices_cmp_global, + ) + else: + bwd_q = q_indexer_flat + bwd_w = w_indexer + bwd_k = k_indexer_flat + bwd_attn = attn_score_for_bwd + bwd_idx = index_score_for_bwd + bwd_topk = topk_indices_cmp + + ig = _DSA.indexer_backward_wrapper( + bwd_q, + bwd_w, + bwd_k, + bwd_attn, + bwd_idx, + bwd_topk, + sm_scale=indexer_softmax_scale, + loss_coeff=indexer_loss_coeff, + grad_loss=unit_grad_loss, + block_I=128, + ) + + if is_thd: + precomputed_grad_q_indexer = ig["d_index_q"].squeeze(0) + precomputed_grad_k_indexer = ig["d_index_k"].squeeze(0) + precomputed_grad_weights = ig["d_weights"].squeeze(0) + else: + precomputed_grad_q_indexer = ig["d_index_q"].permute(1, 0, 2, 3).contiguous() + precomputed_grad_k_indexer = ig["d_index_k"].permute(1, 0, 2).contiguous() + precomputed_grad_weights = ig["d_weights"].permute(1, 0, 2).contiguous() + 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( + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_compressed_idx, + max_seqlen_q=int(max_seqlen_q), + max_seqlen_k=int(max_seqlen_compressed_idx), + ) + ig = _DSA.dense_indexer_backward_wrapper( + q_indexer_flat, + w_indexer, + k_indexer_flat, + attn_score_for_bwd, + attn_l1norm, + index_score_for_bwd, + index_lse_for_bwd, + sm_scale=indexer_softmax_scale, + loss_coeff=indexer_loss_coeff, + grad_loss=unit_grad_loss, + ratio=ratio, + block_I=128, + **dense_bwd_kwargs, + ) + if is_thd: + precomputed_grad_q_indexer = ig["d_index_q"] + precomputed_grad_k_indexer = ig["d_index_k"] + precomputed_grad_weights = ig["d_weights"] + else: + precomputed_grad_q_indexer = ig["d_index_q"].permute(1, 0, 2, 3).contiguous() + precomputed_grad_k_indexer = ig["d_index_k"].permute(1, 0, 2).contiguous() + precomputed_grad_weights = ig["d_weights"].permute(1, 0, 2).contiguous() + else: + precomputed_grad_q_indexer = torch.zeros_like(q_indexer) + precomputed_grad_k_indexer = torch.zeros_like(k_indexer) + precomputed_grad_weights = torch.zeros_like(weights) + + # Zero out pre-computed indexer gradients for padding rows so they + # don't contribute to DSAIndexerLossAutoScaler backward. + if padding_row_mask is not None and loss_coeff > 0: + precomputed_grad_q_indexer[padding_row_mask] = 0 + precomputed_grad_weights[padding_row_mask] = 0 + + # ---- 7. Save context (only sparse-attn bwd tensors + indexer grads). - + ctx.save_for_backward( + q_flat, + kv_flat, + attn_sink, + global_idxs, + out_flat, + lse, + precomputed_grad_q_indexer, + precomputed_grad_k_indexer, + precomputed_grad_weights, + ) + ctx.softmax_scale = softmax_scale + ctx.is_thd = is_thd + ctx.np_ = np_ + ctx.d = d + if is_thd: + ctx.total_q = total_q + else: + ctx.sq = sq + ctx.b = b + ctx.skv = skv + + # ---- Output reshape: layout-specific. -------------------------------- + d_v = out_flat.shape[-1] + if is_thd: + output = out_flat.reshape(total_q, np_ * d_v) + else: + output = out_flat.reshape(sq, b, np_, d_v).reshape(sq, b, np_ * d_v) + return output, indexer_loss + + @staticmethod + def backward(ctx, grad_output, grad_loss): + """Backward: sparse attention bwd + scale pre-computed indexer grads.""" + ( + q_flat, + kv_flat, + attn_sink, + global_idxs, + out_flat, + lse, + precomputed_grad_q_indexer, + precomputed_grad_k_indexer, + precomputed_grad_weights, + ) = ctx.saved_tensors + + is_thd = ctx.is_thd + np_, d = ctx.np_, ctx.d + + # ---- 1. Sparse attn backward (flat layout, layout-agnostic). -------- + d_v = out_flat.shape[-1] + if is_thd: + dO_flat = grad_output.reshape(ctx.total_q, np_, d_v) + else: + sq, b, skv = ctx.sq, ctx.b, ctx.skv + dO_flat = grad_output.reshape(sq * b, np_, d_v) + + attn_bwd = _DSA.sparse_attention_backward_wrapper( + q_flat, + kv_flat, + out_flat, + dO_flat, + lse, + attn_sink, + global_idxs, + softmax_scale=ctx.softmax_scale, + topk_length=None, + ) + if is_thd: + grad_query = attn_bwd["dq"] + grad_kv_full = attn_bwd["dkv"] + else: + grad_query = attn_bwd["dq"].reshape(sq, b, np_, d) + grad_kv_full = attn_bwd["dkv"].reshape(skv, b, d) + d_sink = attn_bwd["d_sink"] + + # ---- 2. Scale pre-computed indexer grads by grad_loss. --------------- + grad_q_indexer = precomputed_grad_q_indexer * grad_loss + grad_k_indexer = precomputed_grad_k_indexer * grad_loss + grad_weights = precomputed_grad_weights * grad_loss + + # Grads: query, kv_full, attn_sink, window_idxs, q_indexer, k_indexer, + # weights, indexer_topk, ratio, softmax_scale, indexer_softmax_scale, + # loss_coeff, sparse_loss, kv_offset, calculate_per_token_loss, + # cu_seqlens_q, cu_seqlens_kv, cu_seqlens_kv_full, + # cu_seqlens_compressed_idx, + # max_seqlen_q, max_seqlen_compressed_idx, + # compressed_kv, cu_seqlens_q_unpadded + return ( + grad_query, + grad_kv_full, + d_sink, + None, + grad_q_indexer, + grad_k_indexer, + grad_weights, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + +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")) + dense_lse = _compute_dense_attn_lse( + query.detach(), + compressed_kv.detach().unsqueeze(1), + softmax_scale, + qhead_per_kv_head=np_, + ratio=ratio, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_k, + max_seqlen_kv=max_seqlen_k, + q_causal_offsets=q_causal_offsets, + ) + attn_score, attn_l1norm = _compute_dense_attn_score( + query.detach(), + compressed_kv.detach().unsqueeze(1), + dense_lse, + 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, + attn_sink: Tensor, + window_idxs: Tensor, + q_indexer: Tensor, + k_indexer: Tensor, + weights: Tensor, + indexer_topk: int, + ratio: int, + softmax_scale: float, + indexer_softmax_scale: float = 1.0, + loss_coeff: float = 0.0, + sparse_loss: bool = False, + kv_offset: int = 0, + calculate_per_token_loss: bool = False, + *, + cu_seqlens_q: Optional[Tensor] = None, + cu_seqlens_kv: Optional[Tensor] = None, + cu_seqlens_kv_full: Optional[Tensor] = None, + cu_seqlens_compressed_idx: Optional[Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_compressed_idx: Optional[int] = None, + compressed_kv: Optional[Tensor] = None, + cu_seqlens_q_unpadded: Optional[Tensor] = None, +) -> Tuple[Tensor, Tensor]: + """Path B (training): fused indexer (+KL loss) + sparse attention. + + Layout is selected by ``cu_seqlens_q``: + + * **SBHD** (``cu_seqlens_q is None``, default): inputs carry an + explicit batch axis; the THD kwargs are ignored. + * **THD packed** (``cu_seqlens_q`` supplied): all four + ``cu_seqlens_*`` and four ``max_seqlen_*`` must be supplied (see + below). Both ``sparse_loss=True`` and ``sparse_loss=False`` are + supported — the sparse-loss path globalizes the per-segment-local + topk indices via ``local_to_global_flat`` and the cuDNN + sparse-indexer-backward kernel addresses K/dK by flat ids. + + See :class:`FusedIndexerSparseAttnFunc` for the detailed data flow. + + SBHD args: + query: ``(sq, b, np, d)`` bf16 SBHD — attention query. + kv_full: ``(skv, b, d)`` bf16 SBD — original + compressed KV. + window_idxs: ``(b, sq, win_topk)`` int32 — local window indices. + q_indexer: ``(sq, b, idx_nh, idx_hd)`` bf16 — indexer query. + k_indexer: ``(n_comp, b, idx_hd)`` bf16 — indexer key (compressed). + weights: ``(sq, b, idx_nh)`` bf16 — raw indexer weights. + kv_offset: start of compressed region within ``kv_full`` (== sq). + + SBHD return: ``output (sq, b, np * d_v)`` + scalar ``indexer_loss``. + + THD args (when ``cu_seqlens_q is not None``): + query: ``(total_q, np, d)`` bf16 — flat-packed Q. + kv_full: ``(total_kv_full, d)`` bf16 — per-segment-concat'd + ``[kv, compressed_kv]`` (built by + :func:`csa.cat_per_segment`). + window_idxs: ``(total_q, win_topk)`` int32 — local-per-segment + window indices. + q_indexer: ``(total_q, idx_nh, idx_hd)`` bf16. + k_indexer: ``(total_comp_idx, idx_hd)`` bf16 — compressed-only K + (== Compressor's output, packed flat). + weights: ``(total_q, idx_nh)`` bf16 — raw. + kv_offset: ignored. + cu_seqlens_q: ``(B+1,)`` int32 CUDA. + cu_seqlens_kv: ``(B+1,)`` int32 — original-KV cu_seqlens. + cu_seqlens_kv_full: ``(B+1,)`` int32 — built by + :func:`csa.build_cu_seqlens_kv_full`. + cu_seqlens_compressed_idx: ``(B+1,)`` int32 — Compressor's + second return value. + max_seqlen_q / max_seqlen_compressed_idx: + per-batch maxima for tile sizing. + + THD return: ``output (total_q, np * d_v)`` + scalar ``indexer_loss``. + + Common args: + attn_sink: ``(np,)`` f32 — learnable sink per head. + indexer_topk: number of top-K compressed positions to select. + ratio: compression ratio used for the causal mask. + softmax_scale: attention ``Q @ K^T`` scale, typically + ``1/sqrt(v_head_dim)``. + indexer_softmax_scale: indexer ``Q @ K^T`` scale, typically + ``1/sqrt(idx_hd)``. Applied internally — caller passes raw + (unscaled) ``weights``. + loss_coeff: coefficient scaling the KL divergence loss. + sparse_loss: if ``True``, KL is computed only over the top-K + positions (cheap); if ``False`` (the default, + matches ``transformer_config.dsa_indexer_use_sparse_loss``), + KL is computed over the full causally-valid KV. See + :class:`FusedIndexerSparseAttnFunc` for the full data flow. + compressed_kv: THD only (required) — ``(total_compressed_kv, d)`` + bf16, the pre-packed compressed KV from the Compressor. Used + by the loss path; THD ``kv_full`` is per-segment concatenated + so it cannot be sliced uniformly the way SBHD ``kv_full`` is. + calculate_per_token_loss: if True, report raw local KL sum and + compensate the cuDNN backward wrappers' local averaging. + cu_seqlens_q_unpadded: THD only (optional) — ``(B+1,)`` int32, + the *unpadded* cumulative Q sequence lengths. When CUDA-graph + padding makes ``cu_seqlens_q`` cover all ``total_q`` rows + (including padding), this tensor supplies the true boundaries + so padding rows are excluded from the indexer KL loss and + backward gradients. Ignored when ``None`` or when it equals + ``cu_seqlens_q``. + """ + if cu_seqlens_q is not None: + missing = [ + name + for name, val in ( + ("cu_seqlens_kv", cu_seqlens_kv), + ("cu_seqlens_kv_full", cu_seqlens_kv_full), + ("cu_seqlens_compressed_idx", cu_seqlens_compressed_idx), + ("max_seqlen_q", max_seqlen_q), + ("max_seqlen_compressed_idx", max_seqlen_compressed_idx), + ("compressed_kv", compressed_kv), + ) + if val is None + ] + if missing: + raise ValueError( + f"fused_indexer_sparse_attn THD mode requires {missing} " "to all be supplied." + ) + return FusedIndexerSparseAttnFunc.apply( + query, + kv_full, + attn_sink, + window_idxs, + q_indexer, + k_indexer, + weights, + indexer_topk, + ratio, + softmax_scale, + indexer_softmax_scale, + loss_coeff, + sparse_loss, + kv_offset, + calculate_per_token_loss, + cu_seqlens_q, + cu_seqlens_kv, + cu_seqlens_kv_full, + cu_seqlens_compressed_idx, + max_seqlen_q, + max_seqlen_compressed_idx, + compressed_kv, + cu_seqlens_q_unpadded, + ) + + +__all__ = [ + "batch_of_row", + "build_flat_topk_idxs", + "local_to_global_flat", + "dsa_sparse_attn", + "indexer_topk", + "fused_indexer_sparse_attn", +] 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 new file mode 100644 index 00000000000..cc975d4ce7e --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -0,0 +1,932 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + +from dataclasses import dataclass +from typing import NoReturn, Optional, Union + +import torch + +from megatron.core import tensor_parallel +from megatron.core.extensions.transformer_engine import HAVE_TE +from megatron.core.models.common.embeddings import ( + RotaryEmbedding, + YarnRotaryEmbedding, + apply_rotary_pos_emb, +) +from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + FineGrainedActivationOffloadingInterface as off_interface, +) +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: + (TEColumnParallelLinear, TELinear, set_save_original_input) = (None, None, None) + + +@torch.compile +def _q_rms_norm(q: torch.Tensor, eps: float) -> torch.Tensor: + """Fused RMS normalization for query tensor (no learnable weight).""" + return q * torch.rsqrt(q.square().mean(-1, keepdim=True) + eps) + + +@dataclass +class DSv4HybridSelfAttentionSubmodules: + """Submodules for the DSv4HybridAttention layer.""" + + q_layernorm: LayerNormBuilder + kv_layernorm: LayerNormBuilder + + linear_q_down_proj: Union[ModuleSpec, type] = None + linear_q_up_proj: Union[ModuleSpec, type] = None + linear_kv_proj: Union[ModuleSpec, type] = None + core_attention: Union[ModuleSpec, type] = None + linear_proj: Union[ModuleSpec, type] = None + + +class DSv4HybridAttention(Attention): + """DeepSeek-v4 Hybrid Attention layer.""" + + def __init__( + self, + config: MLATransformerConfig, + submodules: DSv4HybridSelfAttentionSubmodules, + layer_number: int, + attn_mask_type: AttnMaskType, + attention_type: str, + cp_comm_type: Optional[str] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, + compress_ratio: Optional[int] = None, + name: str | None = None, + ) -> None: + + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + attention_type=attention_type, + attn_mask_type=attn_mask_type, + cp_comm_type=cp_comm_type, + pg_collection=pg_collection, + pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, + name=name, + ) + self.config: MLATransformerConfig + + assert ( + get_pg_size(self.pg_collection.tp) == 1 + ), "DSv4 Hybrid Attention only supports TP size 1." + + assert ( + not self.checkpoint_core_attention + ), "Checkpoint core attention is not supported in DSv4 Hybrid Attention." + assert ( + not self.offload_qkv_linear + ), "Offload qkv linear is not supported in DSv4 Hybrid Attention." + + self.query_projection_size = self.config.v_head_dim * self.config.num_attention_heads + + self.q_head_dim = self.config.v_head_dim + + self.key_hidden_size = self.q_head_dim + self.val_hidden_size = self.config.v_head_dim + + self.recompute_up_proj = ( + self.config.recompute_granularity == 'selective' + and "mla_up_proj" in self.config.recompute_modules + ) + self.qkv_up_checkpoint = None + + self.softmax_scale = None + + # Hybrid C/H/W symbols provide a fixed ratio through their ModuleSpec. The array-driven + # D path falls back to the global per-layer ratio list. + ratio_idx = self.config.num_layers + layer_number - 1 if is_mtp_layer else layer_number - 1 + if compress_ratio is None: + compress_ratio = self.config.csa_compress_ratios[ratio_idx] + use_compressed_yarn = compress_ratio > 1 + rope_base = ( + self.config.csa_compress_rotary_base if use_compressed_yarn else self.config.rotary_base + ) + self._dsv4_compress_ratio = compress_ratio + self._dsv4_rope_base = rope_base + self._dsv4_uses_yarn_rope = use_compressed_yarn + if not use_compressed_yarn: + self.rotary_pos_emb = RotaryEmbedding( + self.config.qk_pos_emb_head_dim, + rotary_percent=self.config.rotary_percent, + rotary_base=rope_base, + cp_group=self.pg_collection.cp, + ) + else: + self.rotary_pos_emb = YarnRotaryEmbedding( + self.config.qk_pos_emb_head_dim, + rotary_base=rope_base, + scaling_factor=self.config.rotary_scaling_factor, + original_max_position_embeddings=self.config.original_max_position_embeddings, + beta_fast=self.config.beta_fast, + beta_slow=self.config.beta_slow, + mscale=self.config.mscale, + mscale_all_dim=self.config.mscale_all_dim, + cp_group=self.pg_collection.cp, + ) + + core_attn_extra_kwargs = { + "rotary_pos_emb": self.rotary_pos_emb, + "compress_ratio": compress_ratio, + "is_mtp_layer": is_mtp_layer, + "name": (name + ".core_attention") if name is not None else None, + } + self.core_attention = build_module( + submodules.core_attention, + config=self.config, + layer_number=self.layer_number, + attn_mask_type=self.attn_mask_type, + attention_type=self.attention_type, + softmax_scale=self.softmax_scale, + k_channels=self.q_head_dim, + v_channels=self.config.v_head_dim, + cp_comm_type=cp_comm_type, + pg_collection=self.pg_collection, + **core_attn_extra_kwargs, + ) + + # Output. + self.o_local_groups = self.config.o_groups + assert ( + self.query_projection_size % self.config.o_groups == 0 + ), "num_attention_heads * v_head_dim must be divisible by o_groups" + group_proj_in_size = self.query_projection_size // self.config.o_groups + group_proj_out_size = self.config.o_groups * self.config.o_lora_rank + + _linear_o_group_proj = torch.empty( + group_proj_out_size, + group_proj_in_size, + device=torch.cuda.current_device(), + dtype=self.config.params_dtype, + ) + self.config.init_method(_linear_o_group_proj) + self.linear_o_group_proj = torch.nn.Parameter(_linear_o_group_proj) + + linear_proj_in_size = self.config.o_groups * self.config.o_lora_rank + + self.linear_proj = build_module( + submodules.linear_proj, + linear_proj_in_size, + self.config.hidden_size, + config=self.config, + init_method=self.config.output_layer_init_method, + bias=self.config.add_bias_linear, + input_is_parallel=True, + skip_bias_add=True, + is_expert=False, + tp_comm_buffer_name='proj', + tp_group=self.pg_collection.tp, + ) + + if ( + HAVE_TE + and isinstance(self.linear_proj, TELinear) + and ( + ( + self.config.fp8 + and self.config.fp8_recipe != 'delayed' + and is_te_min_version("2.6.0dev0") + ) + or (self.config.fp4 and is_te_min_version("2.7.0.dev0")) + ) + ): + # For fp8/fp4 training, the output of the fused core_attn is saved by itself, and + # linear_proj also saves the quantized tensor of this output. Here we set the + # linear_proj to save the original input tensors to avoid the extra memory usage of + # the quantized tensor. + set_save_original_input(self.linear_proj) + + def forward( + self, + hidden_states, + attention_mask, + key_value_states=None, + inference_context=None, + rotary_pos_emb=None, + rotary_pos_cos=None, + rotary_pos_sin=None, + rotary_pos_cos_sin=None, + attention_bias=None, + packed_seq_params=None, + position_ids=None, + sequence_len_offset=None, + *, + inference_params=None, + ): + """Forward pass for DeepSeek-v4 Hybrid Attention""" + assert ( + rotary_pos_emb is None + ), "Rotary position embeddings should not be passed into DSv4HybridAttention." + assert ( + attention_bias is None + ), "Attention bias should not be passed into DSv4HybridAttention." + assert ( + rotary_pos_cos is None and rotary_pos_sin is None + ), "DSv4HybridAttention does not support Flash Decoding" + assert ( + not rotary_pos_cos_sin + ), "Flash-infer rope has not been tested with DSv4HybridAttention." + assert ( + inference_context is None and inference_params is None + ), "Inference is not supported for DSv4HybridAttention." + + if packed_seq_params is not None and packed_seq_params.local_cp_size is not None: + raise ValueError( + "DSv4HybridAttention does not support per-microbatch context-parallel groups." + ) + + cp_group = self.pg_collection.cp + 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.") + + 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, + cp_group, + ) + + # ===================== + # Query, Key, and Value + # ===================== + # Get the query, key and value tensors based on the type of attention - + # self or cross attn. + 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() + key = key.contiguous() + value = value.contiguous() + + # ================================== + # core attention computation + # ================================== + # Need corresponding TE change + core_attn_manager = off_interface( + self.offload_core_attention and self.training, query, "core_attn" + ) + with core_attn_manager as query: + core_attn_out = self.core_attention( + query, + key, + value, + attention_mask, + 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=forced_released_tensors + ) + + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + # reshape to same output shape as unpacked case + # (t, np, hn) -> (t, b=1, h=np*hn) + # t is the pack size = sum (sq_i) + # note that batch is a dummy dimension in the packed case + core_attn_out = core_attn_out.reshape(core_attn_out.size(0), 1, -1) + + if self.recompute_up_proj: + assert self.qkv_up_checkpoint is not None + self.qkv_up_checkpoint.discard_output_and_register_recompute(core_attn_out) + self.qkv_up_checkpoint = None + + # inverse RoPE on last qk_pos_emb_head_dim of each head + seq_len = core_attn_out.size(0) + n_heads = self.num_attention_heads_per_partition + pos_dim = self.config.qk_pos_emb_head_dim + nope_dim = self.config.v_head_dim - pos_dim + core_attn_out = core_attn_out.view(seq_len, core_attn_out.size(1), n_heads, -1) + packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + if packed_seq: + cu_seqlens_kv = ( + packed_seq_params.cu_seqlens_kv_padded + if packed_seq_params.cu_seqlens_kv_padded is not None + else packed_seq_params.cu_seqlens_kv + ) + rope_seqlen = packed_seq_params.max_seqlen_kv + rope_max_seqlen_kv = packed_seq_params.max_seqlen_kv + else: + cu_seqlens_kv = None + rope_seqlen = seq_len + rope_max_seqlen_kv = None + # DSv4 reference (DS-Inf) RoPE is pure rotation (norm-preserving). Yarn's + # concentration factor (mscale) is NOT part of the DSv4 model contract -- + # the model relies on Q/KV RMS-norm + unit-magnitude rotation. Force 1.0. + mscale = 1.0 + rotary_pos_cos = None + rotary_pos_sin = None + if self.config.apply_rope_fusion: + # ``mscale=1.0`` strips yarn's concentration factor from the + # cached cos/sin so the fused kernel matches the unfused + # path's forced ``mscale=1.0`` (DSv4 "pure rotation"). + rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( + rope_seqlen, dtype=hidden_states.dtype, packed_seq=packed_seq, mscale=mscale + ) + rotary_pos_emb = None + assert inference_context is None, "Inference with MLA RoPE fusion is not supported" + assert ( + fused_mla_rope_inplace is not None + ), "Fused MLA RoPE apply is not imported successfully" + elif self._dsv4_uses_yarn_rope: + rotary_pos_emb, _ = self.rotary_pos_emb(rope_seqlen, packed_seq=packed_seq) + else: + rotary_pos_emb = self.rotary_pos_emb(rope_seqlen, packed_seq=packed_seq) + if self.config.apply_rope_fusion: + if use_thd_cp: + global_start = cp_group.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. + assert fused_mla_rope_out_of_place is not None + core_attn_out = fused_mla_rope_out_of_place( + core_attn_out, + rotary_pos_cos, + rotary_pos_sin, + nope_dim, + pos_dim, + cu_seqlens_kv, + cp_group.rank(), + cp_group.size(), + inverse=True, + remove_interleaving=True, + ) + if packed_seq: + core_attn_out = core_attn_out.unsqueeze(1) + elif use_thd_cp: + global_start = cp_group.rank() * core_attn_out.shape[0] + core_attn_out = cp_utils.apply_thd_cp_local_rope_unfused( + core_attn_out, + rotary_pos_emb, + nope_dim, + pos_dim, + cu_seqlens_kv, + global_start, + self.config, + inverse=True, + ) + else: + content_part, rot_part = torch.split( + core_attn_out, [core_attn_out.size(-1) - pos_dim, pos_dim], dim=-1 + ) + # ``_apply_rotary_pos_emb_thd`` documents 3-D ``(total, h, d)`` input + # and adds its own batch dim internally; drop the dummy ``b=1`` axis + # for THD before the rope and add it back after. + if packed_seq: + rot_part_in = rot_part.squeeze(1) + else: + rot_part_in = rot_part + rot_part_out = apply_rotary_pos_emb( + rot_part_in, + rotary_pos_emb, + self.config, + cu_seqlens=cu_seqlens_kv, + mscale=mscale, + cp_group=cp_group, + mla_rotary_interleaved=True, + inverse=True, + mla_output_remove_interleaving=True, + max_seqlen=rope_max_seqlen_kv, + ) + if packed_seq: + rot_part = rot_part_out.unsqueeze(1) + else: + rot_part = rot_part_out + core_attn_out = torch.cat([content_part, rot_part], dim=-1) + core_attn_out = core_attn_out.view(seq_len, core_attn_out.size(1), -1) + + # Grouped output + core_attn_out = core_attn_out.view( + core_attn_out.size(0), core_attn_out.size(1), self.o_local_groups, -1 + ) + wo_a_weight = self.linear_o_group_proj.view( + self.o_local_groups, self.config.o_lora_rank, -1 + ) + core_attn_out = torch.einsum("...gd,grd->...gr", core_attn_out, wo_a_weight) + core_attn_out = core_attn_out.reshape(*core_attn_out.shape[:-2], -1) + + # ================= + # Output. [sq, b, h] + # ================= + attn_proj_manager = off_interface(self.offload_attn_proj, core_attn_out, "attn_proj") + with attn_proj_manager as core_attn_out: + output, bias = self.linear_proj(core_attn_out) + output = attn_proj_manager.group_offload(output, forced_released_tensors=[core_attn_out]) + + return output, bias + + +class DSv4HybridSelfAttention(DSv4HybridAttention): + """DSv4Hybrid Self-attention layer class + + Self-attention layer takes input with size [s, b, h] + and returns output of the same size. + """ + + def __init__( + self, + config: MLATransformerConfig, + submodules: DSv4HybridSelfAttentionSubmodules, + layer_number: int, + attn_mask_type=AttnMaskType.padding, + cp_comm_type: Optional[str] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, + compress_ratio: Optional[int] = None, + name: str | None = None, + ): + if pg_collection is None: + # Compatibility fallback for callers not yet passing process groups explicitly. + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + attn_mask_type=attn_mask_type, + attention_type="self", + cp_comm_type=cp_comm_type, + pg_collection=pg_collection, + pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, + compress_ratio=compress_ratio, + name=name, + ) + + q_down_proj_kwargs = {} + if submodules.linear_q_down_proj in [TELinear]: + q_down_proj_kwargs['parallel_mode'] = 'duplicated' + else: + raise ValueError(f"Unsupported linear_q_down_proj: {submodules.linear_q_down_proj}") + + self.linear_q_down_proj = build_module( + submodules.linear_q_down_proj, + self.config.hidden_size, + self.config.q_lora_rank, + config=self.config, + init_method=self.config.init_method, + bias=False, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name='q_down_proj', + skip_weight_param_allocation=False, + tp_group=None, + name=(name + ".linear_q_down_proj") if name is not None else None, + **q_down_proj_kwargs, + ) + + self.linear_q_up_proj = build_module( + submodules.linear_q_up_proj, + self.config.q_lora_rank, + self.config.num_attention_heads * self.q_head_dim, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=False, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name='q_up_proj', + tp_group=pg_collection.tp, + name=(name + ".linear_q_up_proj") if name is not None else None, + ) + + self.linear_kv_proj = build_module( + submodules.linear_kv_proj, + self.config.hidden_size, + self.config.v_head_dim, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=False, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name='kv_up_proj', + tp_group=pg_collection.tp, + name=(name + ".linear_kv_proj") if name is not None else None, + ) + self.kv_layernorm = submodules.kv_layernorm( + hidden_size=self.config.v_head_dim, + config=self.config, + eps=self.config.layernorm_epsilon, + ) + + self.q_layernorm = submodules.q_layernorm( + hidden_size=self.config.q_lora_rank, + config=self.config, + eps=self.config.layernorm_epsilon, + ) + + def get_query_key_value_tensors( + self, + hidden_states, + key_value_states=None, + position_ids=None, + packed_seq_params=None, + 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] + assert ( + hidden_states.ndim == 3 + ), f"hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D" + assert ( + inference_context is None and inference_params is None + ), "Inference is not supported for DSv4HybridSelfAttention." + + # ========================================= + # Prepare RoPE and seqlen related params + # ========================================= + rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( + inference_context, None, hidden_states, self.config, packed_seq_params + ) + + # rotary_pos_emb:[s, b, 1, 64] + # DSv4 reference (DS-Inf) RoPE is pure rotation (norm-preserving). Yarn's + # concentration factor (mscale) is NOT part of the DSv4 model contract -- + # the model relies on Q/KV RMS-norm + unit-magnitude rotation. Force 1.0. + mscale = 1.0 + rotary_pos_cos = None + rotary_pos_sin = None + packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + if self.config.apply_rope_fusion: + # ``mscale=1.0`` strips yarn's concentration factor from the + # cached cos/sin so the fused kernel matches the unfused + # path's forced ``mscale=1.0`` (DSv4 "pure rotation"). + rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( + rotary_seq_len, dtype=hidden_states.dtype, packed_seq=packed_seq, mscale=mscale + ) + rotary_pos_emb = None + assert inference_context is None, "Inference with MLA RoPE fusion is not supported" + assert ( + fused_mla_rope_inplace is not None + ), "Fused MLA RoPE apply is not imported successfully" + elif self._dsv4_uses_yarn_rope: + rotary_pos_emb, _ = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq) + else: + rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq) + + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + if packed_seq_params.cu_seqlens_q_padded is not None: + cu_seqlens_q = packed_seq_params.cu_seqlens_q_padded + else: + cu_seqlens_q = packed_seq_params.cu_seqlens_q + if packed_seq_params.cu_seqlens_kv_padded is not None: + cu_seqlens_kv = packed_seq_params.cu_seqlens_kv_padded + else: + cu_seqlens_kv = packed_seq_params.cu_seqlens_kv + rope_max_seqlen_q = packed_seq_params.max_seqlen_q + rope_max_seqlen_kv = packed_seq_params.max_seqlen_kv + else: + cu_seqlens_q = cu_seqlens_kv = None + rope_max_seqlen_q = rope_max_seqlen_kv = None + + # ========================================= + # QKV down projection and layernorm + # ========================================= + # 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. + # In Megatron-Core, the qkv shape is [t, 1, h, d]. + # 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 + # ========================================= + + if self.config.q_lora_rank is not None: + # q_compressed: [num_tokens, q_lora_rank] + q_compressed = apply_module(self.q_layernorm)(q_compressed) + + # ========================================= + # QKV up projection and RoPE apply + # ========================================= + + 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, ...]; + otherwise, they maintain the unpacked shape [s, b, ...]. In subsequent code comments, + we uniformly use [num_tokens, ...] to denote [s, b, ...] or [t, ...] for two cases. + """ + # q_compressed: [num_tokens, q_lora_rank] + # q: [num_tokens, n * (qk_head_dim + qk_pos_emb_head_dim)] + q, _ = self.linear_q_up_proj(q_compressed) + + # q: [num_tokens, n, q_head_dim] + 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) + + 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: + 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: + 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] + if packed_seq_params is None: + # Keep direct SBHD forward calls with shorter sequences aligned to + # their inputs. THD reuses the max-length table per packed segment. + 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=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() + + 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) + 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: + 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, + ) + + 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""" + self._backward_kv_proj() + self._backward_q_proj() + self._backward_output_proj() + + def _backward_kv_proj(self): + """Computes weight gradients of KV projection layers""" + self.linear_kv_proj.backward_dw() + + def _backward_q_proj(self): + """Computes weight gradients of Q projection layers""" + self.linear_q_down_proj.backward_dw() + self.linear_q_up_proj.backward_dw() + + def _backward_output_proj(self): + """Computes weight gradients of output projection layer""" + self.linear_proj.backward_dw() + + def set_for_recompute_input_layernorm(self): + """Set the attention layer for recompute input_layernorm. Only needed for fp8/fp4.""" + set_save_original_input(self.linear_q_down_proj) + set_save_original_input(self.linear_kv_proj) diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index bfead2a25c5..b76819605f6 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -3,11 +3,10 @@ import copy import math from dataclasses import dataclass -from typing import Optional, Tuple, Union +from typing import List, Optional, Tuple, Union import torch -from megatron.core import parallel_state from megatron.core.models.common.embeddings import ( RotaryEmbedding, YarnRotaryEmbedding, @@ -296,32 +295,72 @@ def save_loss_to_tracker( return tracker = DSAIndexerLossLoggingHelper.tracker + # Hybrid MTP layer numbers can exceed ``num_layers + mtp_num_layers`` + # because every prediction depth can contain multiple hybrid layers. + needed = max(num_layers, layer_number) if "values" not in tracker: - tracker["values"] = torch.zeros(num_layers, device=torch.cuda.current_device()) + tracker["values"] = torch.zeros(needed, device=torch.cuda.current_device()) + elif tracker["values"].shape[0] < needed: + grown = torch.zeros( + needed, device=tracker["values"].device, dtype=tracker["values"].dtype + ) + grown[: tracker["values"].shape[0]] = tracker["values"] + tracker["values"] = grown tracker["values"][layer_number - 1] += loss.detach() tracker["reduce_group"] = reduce_group 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(): - """Collect and reduce the indexer losses across ranks.""" + def reduce_loss_in_tracker( + pg_collection: ProcessGroupCollection, num_layers: Optional[int] = None + ): + """Collect and reduce indexer losses across every pipeline rank. + + Args: + pg_collection: Process groups used for pipeline and data-parallel reductions. + num_layers: Total number of decoder and MTP layers. When provided, ranks without + local indexer losses contribute zeros to the pipeline-wide reduction. + """ tracker = DSAIndexerLossLoggingHelper.tracker - if "values" not in tracker: + pp_group = pg_collection.pp + + # Pipeline ranks can own different attention variants, so first agree on + # a common tracker size. Cache the result because layer allocation is + # static and the negotiation requires a device-to-host synchronization. + if tracker.get("agreed_size") is not None: + size = tracker["agreed_size"] + else: + local_size = tracker["values"].shape[0] if "values" in tracker else (num_layers or 0) + size_t = torch.tensor( + [local_size], device=torch.cuda.current_device(), dtype=torch.long + ) + torch.distributed.all_reduce(size_t, op=torch.distributed.ReduceOp.MAX, group=pp_group) + size = int(size_t.item()) + tracker["agreed_size"] = size + if size == 0: return + if "values" not in tracker: + tracker["values"] = torch.zeros(size, device=torch.cuda.current_device()) + elif tracker["values"].shape[0] < size: + grown = torch.zeros( + size, device=tracker["values"].device, dtype=tracker["values"].dtype + ) + grown[: tracker["values"].shape[0]] = tracker["values"] + tracker["values"] = grown values = tracker["values"] - torch.distributed.all_reduce( - values, group=parallel_state.get_pipeline_model_parallel_group() - ) + torch.distributed.all_reduce(values, group=pp_group) # Reduce indexer losses across ranks. if tracker.get('reduce_group') is not None: torch.distributed.all_reduce(values, group=tracker.get('reduce_group')) @@ -331,7 +370,7 @@ def reduce_loss_in_tracker(): ) torch.distributed.all_reduce( values, - group=parallel_state.get_data_parallel_group(with_context_parallel=False), + group=pg_collection.dp, op=torch.distributed.ReduceOp.AVG, ) @@ -340,9 +379,13 @@ def track_indexer_metrics( loss_scale: float, iteration: int, writer, + pg_collection: ProcessGroupCollection, wandb_writer=None, total_loss_dict=None, 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. @@ -350,20 +393,30 @@ def track_indexer_metrics( loss_scale: Scale factor for the loss. iteration: Current training iteration. writer: TensorBoard writer. + pg_collection: Process groups used for pipeline and data-parallel reductions. wandb_writer: Weights & Biases writer. total_loss_dict: Dictionary to accumulate total losses. per_layer_logging: Whether to log per-layer losses. + num_layers: Total number of decoder and MTP layers. Passing it makes ranks + without a local indexer participate in the pipeline reduction. + csa_compress_ratios: Per-layer compression ratios. Ratio 4 is the only + DSv4 attention variant that owns an indexer. + preserve_groups: Keep tracker reduction groups after logging for + Transformer Engine graph replay. """ - DSAIndexerLossLoggingHelper.reduce_loss_in_tracker() + DSAIndexerLossLoggingHelper.reduce_loss_in_tracker( + pg_collection=pg_collection, num_layers=num_layers + ) tracker = DSAIndexerLossLoggingHelper.tracker if "values" not in tracker: return indexer_loss_values = tracker["values"] * loss_scale - num_layers = indexer_loss_values.shape[0] - - # Average across all layers (assuming all layers have sparse attention) - avg_indexer_loss = indexer_loss_values.sum() / num_layers + if csa_compress_ratios is not None: + num_indexer_layers = sum(ratio == 4 for ratio in csa_compress_ratios) + else: + num_indexer_layers = indexer_loss_values.shape[0] + avg_indexer_loss = indexer_loss_values.sum() / max(num_indexer_layers, 1) # Log average loss if total_loss_dict is not None: @@ -378,7 +431,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( @@ -595,7 +648,7 @@ def fused_qk_topk_naive( index_scores = index_scores + mask # ========================================= - # Select top-k indices + # Select top-k indices (over the KV axis) # ========================================= topk_k = min(index_topk, sk) if topk_k > 0: @@ -609,6 +662,86 @@ def fused_qk_topk_naive( return index_scores, topk_indices +def fused_qk_topk_naive_thd( + q: torch.Tensor, # (total_q, idx_nh, idx_hd) + k: torch.Tensor, # (total_k, idx_hd) + weights: torch.Tensor, # (total_q, idx_nh) + index_topk: int, + cu_seqlens_q: torch.Tensor, # (B+1,) int32 + cu_seqlens_kv: torch.Tensor, # (B+1,) int32 — indexer-K cu_seqlens + ratio: int, # indexer compression ratio (for causal mask) +): + """THD per-segment naive QK + top-K — the THD analogue of + :func:`fused_qk_topk_naive`. + + For each of the ``B`` segments, slices the per-segment THD inputs + to SBHD with ``b=1``, builds the per-segment compressed-KV causal + mask, delegates to :func:`fused_qk_topk_naive`, and writes the + resulting LOCAL top-K ids back into a flat ``(total_q, index_topk)`` + buffer. Invalid tail positions (rows whose causal-valid count is + smaller than the kernel's top-K width — e.g. early rows with + ``(pos+1)//ratio < index_topk``) are explicitly marked as ``-1`` + so the downstream pipeline can treat them as sentinels (matching + the cuDNN :func:`dsa_kernels.indexer_topk` THD contract). + + This is the unfused code path and the performance is not good. + + Returns: + ``(None, topk_indices_thd)`` where ``topk_indices_thd`` is + ``(total_q, index_topk)`` int64 with per-segment LOCAL ids in + ``[0, seqlen_kv[b])``; ``-1`` for invalid slots. ``index_scores`` + is ``None`` because per-segment scores have heterogeneous + ``(sq_b, sk_b)`` shapes and the only current consumer + (``CompressedSparseAttention._forward_thd`` force_unfused + inference) discards them. + """ + B = int(cu_seqlens_q.shape[0]) - 1 + total_q = q.shape[0] + device = q.device + + topk_thd = torch.full((total_q, index_topk), -1, dtype=torch.int64, device=device) + + for b in range(B): + q_start = int(cu_seqlens_q[b].item()) + q_end = int(cu_seqlens_q[b + 1].item()) + k_start = int(cu_seqlens_kv[b].item()) + k_end = int(cu_seqlens_kv[b + 1].item()) + sq_b = q_end - q_start + sk_b = k_end - k_start + if sq_b == 0 or sk_b == 0: + continue + + # Reshape per-segment to SBHD with b=1; build per-segment mask + # from ratio (same construction as ``_build_causal_mask_seg``). + q_b = q[q_start:q_end].unsqueeze(1) # (sq_b, 1, idx_nh, idx_hd) + k_b = k[k_start:k_end].unsqueeze(1) # (sk_b, 1, idx_hd) + w_b = weights[q_start:q_end].unsqueeze(1) # (sq_b, 1, idx_nh) + mask_b = _build_causal_mask_seg(sq_b, sk_b, ratio, device) + + _, topk_b = fused_qk_topk_naive(q_b, k_b, w_b, index_topk, mask_b) + # topk_b: (1, sq_b, topk_k) where topk_k = min(index_topk, sk_b). + topk_b = topk_b.squeeze(0) + topk_k = topk_b.shape[-1] + + # Mark invalid tail positions per row as ``-1``. A row at + # position ``i`` (0-indexed within the segment) has at most + # ``(i+1) // ratio`` causally-valid compressed positions; any + # topk-slot beyond that count was a ``-inf``-masked selection + # whose value is undefined — convert to the sentinel ``-1`` so + # downstream consumers can ignore it uniformly with the cuDNN + # ``indexer_topk`` contract. + pos_in_seg = torch.arange(sq_b, device=device) + n_valid_per_row = ((pos_in_seg + 1) // ratio).clamp(max=sk_b).clamp(max=topk_k) # (sq_b,) + col_idx = torch.arange(topk_k, device=device).unsqueeze(0) # (1, topk_k) + invalid = col_idx >= n_valid_per_row.unsqueeze(1) # (sq_b, topk_k) + topk_b = torch.where(invalid, torch.full_like(topk_b, -1), topk_b) + + topk_thd[q_start:q_end, :topk_k] = topk_b + # Tail columns [topk_k:index_topk] stay -1 (preallocated full(-1)). + + return None, topk_thd + + def fwd_fused_indexer_loss_naive( q, weights, @@ -872,6 +1005,229 @@ def bwd_fused_indexer_loss_naive( return grad_q.to(q.dtype), grad_weights.to(weights.dtype), grad_k.to(k.dtype) +def _build_causal_mask_seg(seqlen_q_b: int, seqlen_k_b: int, ratio: int, device) -> torch.Tensor: + """Per-segment compressed-KV causal mask ``(1, seqlen_q_b, seqlen_k_b)``. + + Mirrors the SBHD caller's construction in ``csa.py``'s + ``force_unfused_dsa`` branch: column ``j`` is valid for query row ``i`` + iff ``j < (i + 1) // ratio`` (the indexer's bottom-right causal mask + against compressed positions). + """ + cols = torch.arange(seqlen_k_b, device=device).unsqueeze(0).expand(seqlen_q_b, -1) + positions = torch.arange(1, seqlen_q_b + 1, device=device).unsqueeze(1) + return torch.where(cols >= positions // ratio, float('-inf'), 0.0).unsqueeze( + 0 + ) # (1, seqlen_q_b, seqlen_k_b) + + +def fwd_fused_indexer_loss_naive_thd( + q, # (total_q, idx_nh, idx_hd) + weights, # (total_q, idx_nh) — already sm-scale-applied by caller + k, # (total_k_idx, idx_hd) + query, # (total_q, np, hn) — attn Q + key, # (total_k_attn, np, hn) — attn K compressed, expanded MQA + topk, + softmax_scale, + loss_coeff, + sparse_loss, + pg_collection, + cu_seqlens_q, # (B+1,) int32 — shared by indexer Q and attn Q + cu_seqlens_compressed_idx, # (B+1,) int32 — indexer K and attn-compressed K cu_seqlens + ratio, # indexer compression ratio + calculate_per_token_loss=False, +): + """THD per-segment forward — loops over segments and delegates each + one to :func:`fwd_fused_indexer_loss_naive` with ``b=1``. + + Returns ``(topk_indices_thd (total_q, topk) int32 [per-segment LOCAL + ids], indexer_loss (scalar))``. Aggregation matches the SBHD + definition for each reduction mode: + + * **mean** (``calculate_per_token_loss=False``): ``loss_b`` is the + per-segment row MEAN, so weight by the segment length and divide by + ``total_q`` to recover the row-mean over ALL THD query rows:: + + ``loss = sum_b (loss_b * seqlen_q[b]) / total_q`` + + * **per-token** (``calculate_per_token_loss=True``): ``loss_b`` is + already a RAW ROW SUM over the segment's rows, so the aggregate is a + plain ``sum_b loss_b`` over all THD rows (the global token divisor is + applied later by ``finalize_model_grads``). The mean-mode + ``* seqlen_q[b] / total_q`` weighting must NOT be applied here — doing + so scales the loss (and every indexer gradient) by ``1 / num_segments``. + + Segments with ``seqlen_k[b] == 0`` contribute nothing (mean-mode still + counts their rows in ``total_q``). + """ + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "fwd_fused_indexer_loss_naive_thd: this unfused per-segment loop uses " + "GPU→CPU syncs (.item()) and cannot run during CUDA graph capture. " + "Use the fused kernel path (apply_dsa_kernel_fusion=True) instead." + ) + + B = int(cu_seqlens_q.shape[0]) - 1 + total_q = q.shape[0] + device = q.device + + topk_indices_thd = torch.full((total_q, topk), -1, dtype=torch.int32, device=device) + weighted_losses = [] + + for b in range(B): + q_start = int(cu_seqlens_q[b].item()) + q_end = int(cu_seqlens_q[b + 1].item()) + k_start = int(cu_seqlens_compressed_idx[b].item()) + k_end = int(cu_seqlens_compressed_idx[b + 1].item()) + seqlen_q_b = q_end - q_start + seqlen_k_b = k_end - k_start + if seqlen_q_b == 0 or seqlen_k_b == 0: + continue + + # Slice per-segment; reshape to SBHD with b=1 (the existing + # naive helpers' contract). Each ``unsqueeze(1)`` is a view — + # the per-segment compute reuses storage from the THD tensors. + q_b = q[q_start:q_end].unsqueeze(1) + weights_b = weights[q_start:q_end].unsqueeze(1) + k_b = k[k_start:k_end].unsqueeze(1) + query_b = query[q_start:q_end].unsqueeze(1) + key_b = key[k_start:k_end].unsqueeze(1) + mask_b = _build_causal_mask_seg(seqlen_q_b, seqlen_k_b, ratio, device) + + topk_indices_b, loss_b = fwd_fused_indexer_loss_naive( + q_b, + weights_b, + k_b, + query_b, + key_b, + topk, + softmax_scale, + loss_coeff, + mask_b, + sparse_loss, + pg_collection, + calculate_per_token_loss=calculate_per_token_loss, + ) + # topk_indices_b: (1, seqlen_q_b, topk_seg) where + # ``topk_seg = min(topk, seqlen_k_b)``. Real segments with + # ``seqlen_k_b < topk`` produce a narrower slice; write only + # those columns and leave the trailing ``[topk_seg:topk]`` + # range at the buffer's initial -1 sentinel so the downstream + # post-filter in csa.py marks them invalid. + topk_seg = topk_indices_b.shape[-1] + topk_indices_thd[q_start:q_end, :topk_seg] = topk_indices_b.squeeze(0).int() + # per-token: ``loss_b`` is a raw row sum -> aggregate is a plain sum. + # mean: ``loss_b`` is a row mean -> weight by segment length here and + # divide by ``total_q`` below to get the row-mean over all THD rows. + weighted_losses.append(loss_b if calculate_per_token_loss else loss_b * seqlen_q_b) + + if weighted_losses: + indexer_loss = torch.stack(weighted_losses).sum() + if not calculate_per_token_loss: + indexer_loss = indexer_loss / float(max(total_q, 1)) + else: + indexer_loss = torch.zeros((), device=device, dtype=torch.float32) + return topk_indices_thd, indexer_loss + + +def bwd_fused_indexer_loss_naive_thd( + q, + weights, + k, + query, + key, + topk_indices_thd, + softmax_scale, + loss_coeff, + sparse_loss, + grad_loss, + pg_collection, + cu_seqlens_q, + cu_seqlens_compressed_idx, + ratio, + calculate_per_token_loss=False, +): + """THD per-segment backward — accumulates per-segment grads back into + the flat THD-shaped grad buffers. + + The per-segment ``grad_loss`` must match the forward's aggregation + (see :func:`fwd_fused_indexer_loss_naive_thd`): + + * **mean** (``calculate_per_token_loss=False``): scale by + ``seqlen_q[b] / total_q`` so the inner naive backward's internal + ``/seqlen_q[b]`` row-mean divisor composes into the correct per-row + gradient of the row-weighted-mean aggregate. + * **per-token** (``calculate_per_token_loss=True``): the aggregate is a + plain ``sum_b loss_b`` and the inner backward does NOT divide, so each + segment carries the FULL upstream ``grad_loss``. Applying the mean-mode + ``seqlen_q[b] / total_q`` factor here would shrink every indexer + gradient by ``1 / num_segments``. + """ + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "bwd_fused_indexer_loss_naive_thd: this unfused per-segment loop uses " + "GPU→CPU syncs (.item()) and cannot run during CUDA graph capture. " + "Use the fused kernel path (apply_dsa_kernel_fusion=True) instead." + ) + + B = int(cu_seqlens_q.shape[0]) - 1 + device = q.device + total_q = max(int(q.shape[0]), 1) + + grad_q = torch.zeros_like(q) + grad_weights = torch.zeros_like(weights) + grad_k = torch.zeros_like(k) + + for b in range(B): + q_start = int(cu_seqlens_q[b].item()) + q_end = int(cu_seqlens_q[b + 1].item()) + k_start = int(cu_seqlens_compressed_idx[b].item()) + k_end = int(cu_seqlens_compressed_idx[b + 1].item()) + seqlen_q_b = q_end - q_start + seqlen_k_b = k_end - k_start + if seqlen_q_b == 0 or seqlen_k_b == 0: + continue + + q_b = q[q_start:q_end].unsqueeze(1) + weights_b = weights[q_start:q_end].unsqueeze(1) + k_b = k[k_start:k_end].unsqueeze(1) + query_b = query[q_start:q_end].unsqueeze(1) + key_b = key[k_start:k_end].unsqueeze(1) + # Slice to ``min(topk_global, seqlen_k_b)`` so segments whose + # K count is shorter than the global topk don't feed -1 + # sentinels (the buffer's initial value) into the inner + # ``bwd_fused_indexer_loss_naive``'s ``scatter_(-1, ..., 0)``, + # which would OOB. The forward writes only this many columns. + topk_seg = min(topk_indices_thd.shape[-1], seqlen_k_b) + topk_b = topk_indices_thd[q_start:q_end, :topk_seg].unsqueeze(0).long() + mask_b = _build_causal_mask_seg(seqlen_q_b, seqlen_k_b, ratio, device) + + # per-token: plain sum aggregate -> full grad per segment. + # mean: scale by (seqlen_q_b / total_q) so the inner naive backward's + # internal /seqlen_q_b divisor yields the row-mean over all THD rows. + grad_loss_b = grad_loss if calculate_per_token_loss else grad_loss * (seqlen_q_b / total_q) + + grad_q_b, grad_w_b, grad_k_b = bwd_fused_indexer_loss_naive( + q_b, + weights_b, + k_b, + query_b, + key_b, + topk_b, + softmax_scale, + loss_coeff, + sparse_loss, + mask_b, + grad_loss_b, + pg_collection, + calculate_per_token_loss=calculate_per_token_loss, + ) + grad_q[q_start:q_end] += grad_q_b.squeeze(1) + grad_weights[q_start:q_end] += grad_w_b.squeeze(1) + grad_k[k_start:k_end] += grad_k_b.squeeze(1) + + return grad_q, grad_weights, grad_k + + _FUSED_DSA_INDEXER_LOSS_INPUT_NAMES = ( "q", "weights", @@ -890,11 +1246,45 @@ def bwd_fused_indexer_loss_naive( "query_valid_rows", "calculate_per_token_loss", "use_relu", + "cu_seqlens_q", + "cu_seqlens_compressed_idx", + "ratio", ) class FusedDSAIndexerLoss(torch.autograd.Function): - """Fused implementation of DSA Indexer Loss.""" + """Fused implementation of DSA Indexer Loss. + + Supports both SBHD (default) and THD packed-sequence layouts. THD + is selected by passing ``cu_seqlens_q`` (and the corresponding + ``cu_seqlens_compressed_idx`` + ``ratio``) — those args are appended + at the end of the positional signature so the existing SBHD callers + remain source-compatible (they pass ``None`` / are unchanged). + + SBHD shapes: + q (sq, b, idx_nh, idx_hd) + weights (sq, b, idx_nh) + k (sk, b, idx_hd) + query (sq, b, np, hn) + key (sk, b, np, hn) (compressed-only, MQA-expanded) + mask (b, sq, sk) — caller-built per-batch causal mask. + + THD shapes (``cu_seqlens_q`` supplied): + q (total_q, idx_nh, idx_hd) + weights (total_q, idx_nh) + k (total_k_idx, idx_hd) + query (total_q, np, hn) + key (total_k_attn, np, hn) (compressed-only, MQA-expanded; + ``total_k_attn == total_k_idx`` because both come from + same-ratio compressors over the same input lengths) + mask ignored — built per-segment internally from ``ratio``. + + Implementation: SBHD uses the existing single-pass naive helpers; + THD loops over segments and delegates each one to the same SBHD + helpers with ``b=1`` (the math is identical per-segment, and the + per-row mean is recovered via a row-weighted average of the + per-segment losses). + """ @staticmethod def forward( @@ -916,31 +1306,58 @@ def forward( query_valid_rows=None, calculate_per_token_loss: bool = False, use_relu: bool = True, + cu_seqlens_q=None, + cu_seqlens_compressed_idx=None, + ratio=None, ): """ Fused forward: index_scores never materialized in full. """ - topk_indices, loss = fwd_fused_indexer_loss_naive( - q, - weights, - k, - query, - key, - topk, - softmax_scale, - loss_coeff, - mask, - sparse_loss, - pg_collection, - varlen_starts=varlen_starts, - varlen_ends=varlen_ends, - key_positions=key_positions, - query_valid_rows=query_valid_rows, - calculate_per_token_loss=calculate_per_token_loss, - use_relu=use_relu, - ) + is_thd = cu_seqlens_q is not None + if is_thd: + if cu_seqlens_compressed_idx is None or ratio is None: + raise ValueError( + "FusedDSAIndexerLoss THD mode requires both " + "``cu_seqlens_compressed_idx`` and ``ratio``." + ) + topk_indices, loss = fwd_fused_indexer_loss_naive_thd( + q, + weights, + k, + query, + key, + topk, + softmax_scale, + loss_coeff, + sparse_loss, + pg_collection, + cu_seqlens_q, + cu_seqlens_compressed_idx, + ratio, + calculate_per_token_loss, + ) + else: + topk_indices, loss = fwd_fused_indexer_loss_naive( + q, + weights, + k, + query, + key, + topk, + softmax_scale, + loss_coeff, + mask, + sparse_loss, + pg_collection, + varlen_starts=varlen_starts, + varlen_ends=varlen_ends, + key_positions=key_positions, + query_valid_rows=query_valid_rows, + calculate_per_token_loss=calculate_per_token_loss, + use_relu=use_relu, + ) - # Save for backward (recomputation strategy) + # THD rebuilds per-segment masks in the backward. ctx.save_for_backward(q, weights, k, query, key, topk_indices) ctx.softmax_scale = softmax_scale ctx.loss_coeff = loss_coeff @@ -953,6 +1370,10 @@ def forward( ctx.query_valid_rows = query_valid_rows ctx.calculate_per_token_loss = calculate_per_token_loss ctx.use_relu = use_relu + ctx.is_thd = is_thd + ctx.cu_seqlens_q = cu_seqlens_q + ctx.cu_seqlens_compressed_idx = cu_seqlens_compressed_idx + ctx.ratio = ratio return topk_indices, loss @@ -963,26 +1384,45 @@ def backward(ctx, grad_topk_indices, grad_loss): """ q, weights, k, query, key, topk_indices = ctx.saved_tensors - grad_q, grad_weights, grad_k = bwd_fused_indexer_loss_naive( - q, - weights, - k, - query, - key, - topk_indices, - ctx.softmax_scale, - ctx.loss_coeff, - ctx.sparse_loss, - ctx.mask, - grad_loss, - ctx.pg_collection, - varlen_starts=ctx.varlen_starts, - varlen_ends=ctx.varlen_ends, - key_positions=ctx.key_positions, - query_valid_rows=ctx.query_valid_rows, - calculate_per_token_loss=ctx.calculate_per_token_loss, - use_relu=ctx.use_relu, - ) + if ctx.is_thd: + grad_q, grad_weights, grad_k = bwd_fused_indexer_loss_naive_thd( + q, + weights, + k, + query, + key, + topk_indices, + ctx.softmax_scale, + ctx.loss_coeff, + ctx.sparse_loss, + grad_loss, + ctx.pg_collection, + ctx.cu_seqlens_q, + ctx.cu_seqlens_compressed_idx, + ctx.ratio, + calculate_per_token_loss=ctx.calculate_per_token_loss, + ) + else: + grad_q, grad_weights, grad_k = bwd_fused_indexer_loss_naive( + q, + weights, + k, + query, + key, + topk_indices, + ctx.softmax_scale, + ctx.loss_coeff, + ctx.sparse_loss, + ctx.mask, + grad_loss, + ctx.pg_collection, + varlen_starts=ctx.varlen_starts, + varlen_ends=ctx.varlen_ends, + key_positions=ctx.key_positions, + query_valid_rows=ctx.query_valid_rows, + calculate_per_token_loss=ctx.calculate_per_token_loss, + use_relu=ctx.use_relu, + ) grad_by_name = { "q": grad_q, diff --git a/megatron/core/transformer/experimental_attention_variant/dsv4_module_specs.py b/megatron/core/transformer/experimental_attention_variant/dsv4_module_specs.py new file mode 100644 index 00000000000..e142413283e --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/dsv4_module_specs.py @@ -0,0 +1,95 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Backend-neutral module specs for DeepSeek-V4 compressed sparse attention.""" + +from typing import Protocol + +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.experimental_attention_variant.csa import ( + CompressedSparseAttention, + CompressedSparseAttentionSubmodules, + Compressor, + CompressorSubmodules, + CSAIndexer, + CSAIndexerSubmodules, +) +from megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention import ( + DSv4HybridSelfAttention, + DSv4HybridSelfAttentionSubmodules, +) +from megatron.core.transformer.identity_op import IdentityOp +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_config import TransformerConfig + + +class DSv4BackendSpecProvider(Protocol): + """Minimal backend interface needed to build the DSv4 attention spec.""" + + def linear(self) -> type: + """Return the backend's non-parallel linear module.""" + ... + + def column_parallel_linear(self) -> type: + """Return the backend's column-parallel linear module.""" + ... + + def row_parallel_linear(self) -> type: + """Return the backend's row-parallel linear module.""" + ... + + def layer_norm( + self, rms_norm: bool = False, for_qk: bool = False, has_residual: bool = False + ) -> type: + """Return the backend's normalization module.""" + ... + + +def get_dsv4_hybrid_module_spec_for_backend( + config: TransformerConfig, backend: DSv4BackendSpecProvider +) -> ModuleSpec: + """Build a DSv4 compressed sparse-attention spec for an explicit backend.""" + assert config.multi_latent_attention, "Currently only MLA supports sparse attention." + assert config.qk_l2_norm is False, "qk_l2_norm is not supported with MLA." + + rms_norm = config.normalization == "RMSNorm" + qk_norm = ( + backend.layer_norm(rms_norm=rms_norm, for_qk=True) if config.qk_layernorm else IdentityOp + ) + + compressor_spec = ModuleSpec( + module=Compressor, + submodules=CompressorSubmodules( + linear_wkv=backend.linear(), + linear_wgate=backend.linear(), + norm=backend.layer_norm(rms_norm=True, for_qk=False), + ), + ) + indexer_spec = ModuleSpec( + module=CSAIndexer, + submodules=CSAIndexerSubmodules( + linear_wq_b=backend.linear(), + linear_weights_proj=backend.linear(), + compressor=compressor_spec, + ), + ) + core_attention = ModuleSpec( + module=CompressedSparseAttention, + submodules=CompressedSparseAttentionSubmodules( + compressor=compressor_spec, indexer=indexer_spec + ), + ) + + return ModuleSpec( + module=DSv4HybridSelfAttention, + params={"attn_mask_type": AttnMaskType.causal}, + submodules=DSv4HybridSelfAttentionSubmodules( + linear_q_down_proj=backend.linear(), + linear_q_up_proj=backend.column_parallel_linear(), + linear_kv_proj=backend.column_parallel_linear(), + core_attention=core_attention, + linear_proj=backend.row_parallel_linear(), + q_layernorm=qk_norm, + kv_layernorm=qk_norm, + ), + metainfo={"fuse_input_layernorm": False}, + ) diff --git a/megatron/core/transformer/hyper_connection.py b/megatron/core/transformer/hyper_connection.py new file mode 100644 index 00000000000..44ec8e2144d --- /dev/null +++ b/megatron/core/transformer/hyper_connection.py @@ -0,0 +1,843 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import math +from typing import TYPE_CHECKING, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor + +from megatron.core.transformer.module import MegatronModule, mark_keep_in_fp32 +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import nvtx_decorator + +if TYPE_CHECKING: + from megatron.core.tensor_parallel.random import CheckpointManager + +_MHC_SINKHORN_EPS = 1e-6 +_MHC_COMPUTE_H_EPS = 1e-6 + + +@torch.compile +def _sinkhorn_iterations(input_logits: Tensor, num_iterations: int, eps: float) -> Tensor: + M = input_logits.softmax(dim=-1) + eps + M = M / (M.sum(dim=-2, keepdim=True) + eps) + for _ in range(num_iterations - 1): + M = M / (M.sum(dim=-1, keepdim=True) + eps) + M = M / (M.sum(dim=-2, keepdim=True) + eps) + return M + + +class SinkhornKnopp(torch.autograd.Function): + """Sinkhorn-Knopp projection to doubly stochastic matrix. + + This is an autograd.Function because the iterative forward is re-executed + during backward (under torch.enable_grad) so that PyTorch's autograd can + differentiate through it without storing all intermediate iteration states. + """ + + @staticmethod + def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6) -> Tensor: + """Run Sinkhorn iterations and save inputs for backward recomputation.""" + M = _sinkhorn_iterations(input_logits, num_iterations, eps) + ctx.save_for_backward(input_logits) + ctx.num_iterations = num_iterations + ctx.eps = eps + return M + + @staticmethod + def backward(ctx, grad_output: Tensor): + """Recompute forward under enable_grad and back-propagate.""" + (input_logits,) = ctx.saved_tensors + with torch.enable_grad(): + logits = input_logits.detach().requires_grad_(True) + M = _sinkhorn_iterations(logits, ctx.num_iterations, ctx.eps) + M.backward(grad_output) + return logits.grad, None, None + + +def native_sinkhorn(input_logits: Tensor, num_iterations: int, eps: float = 1e-6) -> Tensor: + """Native Sinkhorn-Knopp (autograd.Function wrapper).""" + return SinkhornKnopp.apply(input_logits, num_iterations, eps) + + +@torch.compile +def native_h_aggregate(x: Tensor, h_pre: Tensor) -> Tensor: + """Native n-stream weighted aggregation: out = sum_j(h_pre_j * x_j).""" + return (x * h_pre.unsqueeze(-1)).sum(dim=2) + + +@torch.compile +def native_h_post_bda( + h_res: Tensor, original_residual: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] +) -> Tensor: + """Native H_res.T @ residual + H_post * (x [+ bias]).""" + s, b, n, C = original_residual.shape + h_res_batched = h_res.view(s * b, n, n) + residual_batched = original_residual.view(s * b, n, C) + mixed = torch.bmm(h_res_batched.transpose(1, 2), residual_batched).view(s, b, n, C) + x_expanded = h_post.unsqueeze(-1) * x.unsqueeze(2) + if bias is not None: + bias_expanded = h_post.unsqueeze(-1) * bias.view(1, 1, 1, C) + return x_expanded + bias_expanded + mixed + return x_expanded + mixed + + +@torch.compile +def native_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6) -> Tuple[Tensor, Tensor]: + """Native fused projection + RMS normalization.""" + proj = torch.matmul(x, weight.t()) + norm = x.norm(dim=-1, keepdim=True) + K = x.shape[-1] + v = norm / math.sqrt(K) + eps + r = 1.0 / v + return proj, r + + +@torch.compile +def native_fused_add_3(a: Tensor, b: Tensor, c: Tensor) -> Tensor: + """Native 3-way elementwise add (torch.compile fuses into single kernel).""" + return a + b + c + + +class BroadcastTensorFused(torch.autograd.Function): + """Split one tensor into 3 autograd-graph children sharing the same storage. + + During backward the three incoming gradients are summed with a caller- + supplied fused-add function (cuTile or torch.compile fallback) instead of + PyTorch's default sequential accumulation. + """ + + @staticmethod + def forward(ctx, x, fused_add_3_fn): + """Return three view aliases and save the fused gradient combiner.""" + ctx.fused_add_3_fn = fused_add_3_fn + return x.view_as(x), x.view_as(x), x.view_as(x) + + @staticmethod + def backward(ctx, grad1, grad2, grad3): + """Combine gradients from the three broadcast aliases.""" + grads = [g for g in (grad1, grad2, grad3) if g is not None] + if len(grads) == 0: + return None, None + if len(grads) == 1: + return grads[0], None + if len(grads) == 2: + return grads[0] + grads[1], None + return ctx.fused_add_3_fn(grad1, grad2, grad3), None + + +@torch.compile +def learned_output_contract( + hidden_states: Tensor, head_fn: Tensor, base: Tensor, scale: Tensor, n: int, eps: float +) -> Tensor: + """Learned output contraction: n-stream → 1-stream via sigmoid-gated weighted sum.""" + dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + head_fn = head_fn.to(torch.float32) + base = base.to(torch.float32) + scale = scale.to(torch.float32) + rsqrt = torch.rsqrt(hidden_states.square().mean(-1, keepdim=True) + eps) + mixes = F.linear(hidden_states, head_fn) * rsqrt + pre = torch.sigmoid(mixes * scale + base) + eps + y = torch.sum(pre.unsqueeze(-1) * hidden_states.view(*hidden_states.shape[:-1], n, -1), dim=-2) + return y.to(dtype) + + +# ============================================================================ +# HyperConnectionModule +# ============================================================================ + + +# TODO: keep hyper connection in fp32 computation +class HyperConnectionModule(MegatronModule): + """ + Unified mHC (Manifold-Constrained Hyper-Connections) module. + + Implements the complete mHC propagation: + x_{l+1} = H_res^T @ x_l + H_post^T @ F(H_pre @ x_l) + + This module handles: + 1. Computing learnable mappings: H_pre, H_post, H_res (with Sinkhorn-Knopp projection) + 2. Aggregation: n-stream → 1-stream (H_pre @ x) + 3. Expansion: 1-stream → n-stream (H_post^T @ output) + 4. Residual merge: H_res^T @ x + expanded_output + 5. Block-level expand/contract for TransformerBlock boundaries + + Args: + config: TransformerConfig with hyper-connection fields + layer_number: Current layer index for initialization + """ + + def __init__(self, config: TransformerConfig, layer_number: int): + super().__init__(config) + self.config = config + self.layer_number = layer_number + self.n = config.num_residual_streams + self.hidden_size = config.hidden_size + self.sinkhorn_iterations = config.mhc_sinkhorn_iterations + self.sinkhorn_eps = _MHC_SINKHORN_EPS + self.compute_h_eps = _MHC_COMPUTE_H_EPS + + # Projection weights for dynamic mappings + # Input: [s, b, n*C] -> Output: n^2 + 2n values per token + # - H_pre: n values + # - H_post: n values + # - H_res: n^2 values (before Sinkhorn projection) + self.mapping_proj = nn.Linear( + self.n * self.hidden_size, self.n * self.n + 2 * self.n, bias=False + ) + + init_alpha = config.mhc_init_gating_factor + # Learnable scaling factors (Eq. 5 in paper) + self.alpha_pre = nn.Parameter(torch.full((1,), init_alpha)) + self.alpha_post = nn.Parameter(torch.full((1,), init_alpha)) + self.alpha_res = nn.Parameter(torch.full((1,), init_alpha)) + + # Static bias terms + self.bias = nn.Parameter(torch.zeros(self.n * self.n + 2 * self.n)) + mark_keep_in_fp32(self.mapping_proj.weight) + mark_keep_in_fp32(self.alpha_pre) + mark_keep_in_fp32(self.alpha_post) + mark_keep_in_fp32(self.alpha_res) + mark_keep_in_fp32(self.bias) + self.norm_eps = 1e-6 + + # Choose implementation: unified fused kernels vs reference modules. + # The fused public API selects the backend per operation internally. + # fused_add_3 always uses torch.compile (native_fused_add_3) regardless + # of the kernel backend. cuTile's register overhead (56 regs/thread for + # a trivial a+b+c) is not worth it for a pure memory-bound elementwise op. + self._fused_add_3_op = native_fused_add_3 + + if config.use_fused_mhc: + from megatron.core.fusions.fused_mhc_kernels import ( + fused_h_aggregate, + fused_h_post_bda, + fused_proj_rms, + fused_proj_rms_compute_h, + fused_sinkhorn, + log_fused_mhc_backend_once, + ) + + log_fused_mhc_backend_once() + self._sinkhorn_op = fused_sinkhorn + self._h_aggregate_op = fused_h_aggregate + self._h_post_bda_op = fused_h_post_bda + self._proj_rms_op = fused_proj_rms + self._proj_rms_compute_h_op = fused_proj_rms_compute_h + else: + self._sinkhorn_op = native_sinkhorn + self._h_aggregate_op = native_h_aggregate + self._h_post_bda_op = native_h_post_bda + self._proj_rms_op = native_proj_rms + self._proj_rms_compute_h_op = None + + self._init_weights() + + def _init_weights(self) -> None: + """Initialize weights for stable training.""" + nn.init.xavier_uniform_(self.mapping_proj.weight) + + # Set sequence_parallel attribute on parameters for gradient synchronization + # across TP ranks when sequence_parallel is enabled. + # This is required because HyperConnectionModule uses non-TP-aware layers + # (nn.Linear, nn.RMSNorm) whose gradients need to be all-reduced. + if self.config.sequence_parallel: + setattr(self.mapping_proj.weight, 'sequence_parallel', True) + setattr(self.alpha_pre, 'sequence_parallel', True) + setattr(self.alpha_post, 'sequence_parallel', True) + setattr(self.alpha_res, 'sequence_parallel', True) + setattr(self.bias, 'sequence_parallel', True) + + def _projection_and_get_norm(self, x: Tensor) -> Tuple[Tensor, Tensor]: + """ + Projection + RMS normalization. + + Args: + x: [s, b, n*C] - n-stream hidden states + """ + s, b, nC = x.shape + # The mHC mapping computation runs in FP32: the parameters are kept in + # FP32 and the activations are upcast here, then compute_mappings casts + # the bounded mixing weights back to the activation dtype. + x_2d = x.reshape(s * b, nC).to(torch.float32) + weight = self.mapping_proj.weight.to(torch.float32) + proj, r = self._proj_rms_op(x_2d, weight, self.norm_eps) + return proj.view(s, b, -1), r.view(s, b, 1) + + @torch.compile + def _compute_h(self, proj: Tensor, r: Tensor) -> Tuple[Tensor, Tensor, Tensor]: + """ + Compute h from projected hidden states and scaling factors. + + Args: + proj: [s, b, n^2 + 2n] - projected hidden states + r: [s, b, 1] - scaling factors + + Returns: + h_pre: [s, b, n] - aggregation weights + h_post: [s, b, n] - expansion weights + h_res: [s, b, n^2] - residual mixing logits + """ + alpha_ = torch.cat( + [ + self.alpha_pre.expand(self.n), + self.alpha_post.expand(self.n), + self.alpha_res.expand(self.n * self.n), + ], + dim=-1, + ) + + h = r * proj * alpha_ + self.bias + # H_pre = σ(α_pre * (θ_pre @ x̃) + b_pre) + h_pre = h[..., : self.n].sigmoid() + self.compute_h_eps # [s, b, n] + + # H_post = 2σ(α_post * (θ_post @ x̃) + b_post) + h_post = h[..., self.n : 2 * self.n].sigmoid() * 2 + h_res = h[..., 2 * self.n :] + return h_pre, h_post, h_res + + @nvtx_decorator(message="HyperConnection::compute_mappings") + def compute_mappings(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor]: + """ + Compute mHC mappings from input hidden states. + + Reference: Eq. (5) and (8) in mHC paper + + Args: + x: [s, b, n*C] - n-stream hidden states + + Returns: + h_pre: [s, b, n] - aggregation weights (sigmoid activated) + h_post: [s, b, n] - expansion weights (2*sigmoid activated) + h_res: [s, b, n, n] - residual mixing matrix (doubly stochastic) + """ + s, b, _ = x.shape + + if self._proj_rms_compute_h_op is not None: + # Fused path: proj_rms + compute_h in one kernel launch sequence + x_2d = x.reshape(s * b, self.n * self.hidden_size) + with torch.cuda.nvtx.range("HyperConnection::fused_proj_rms_compute_h"): + h_pre, h_post, h_res, _ = self._proj_rms_compute_h_op( + x_2d, + self.mapping_proj.weight, + self.alpha_pre, + self.alpha_post, + self.alpha_res, + self.bias, + self.n, + self.norm_eps, + self.compute_h_eps, + ) + h_pre = h_pre.view(s, b, self.n) + h_post = h_post.view(s, b, self.n) + h_res = h_res.view(s, b, self.n, self.n) + else: + # Native path: separate proj_rms + _compute_h + with torch.cuda.nvtx.range("HyperConnection::projection_and_get_norm"): + proj, r = self._projection_and_get_norm(x) + with torch.cuda.nvtx.range("HyperConnection::compute_h"): + h_pre, h_post, h_res = self._compute_h(proj, r) + h_res = h_res.view(s, b, self.n, self.n) + + h_res = self._sinkhorn_op( + h_res, self.sinkhorn_iterations, self.sinkhorn_eps + ) # [s, b, n, n] + + # The mixing weights are bounded (sigmoid outputs / doubly stochastic + # matrix), so after the FP32 computation they are safe to apply to the + # streams in the activation dtype. + dtype = x.dtype + return h_pre.to(dtype), h_post.to(dtype), h_res.to(dtype) + + @torch.compile + def _apply_h_post(self, x: Tensor, h_post: Tensor) -> Tensor: + """ + Core implementation of H_post application to a single tensor. + + Computes: H_post^T @ x + + Args: + x: Input tensor, can be either: + - [s, b, C] - standard hidden states + - [C] - bias tensor (will be broadcast) + h_post: [s, b, n] - expansion weights + + Returns: + output: [s, b, n*C] - expanded tensor + """ + n = self.n + s, b, _ = h_post.shape + + if x.dim() == 1: + # x is bias with shape [C], need to broadcast to [s, b, 1, C] + C = x.shape[0] + x_expanded = x.unsqueeze(0).unsqueeze(0).unsqueeze(0).expand(s, b, 1, C) + else: + # x is [s, b, C] + C = x.shape[-1] + x_expanded = x.unsqueeze(2) # [s, b, 1, C] + + # h_post^T @ x : [s, b, n, 1] * [s, b, 1, C] -> [s, b, n, C] + # Using broadcast multiply instead of einsum + result = h_post.unsqueeze(-1) * x_expanded + return result.view(s, b, n * C) + + @nvtx_decorator(message="HyperConnection::apply_h_post") + def apply_h_post( + self, + x_with_bias: Tuple[Tensor, Optional[Tensor]], + h_post: Tensor, + manager: Optional['CheckpointManager'] = None, + ) -> Tuple[Tensor, Optional[Tensor]]: + """ + Apply H_post to x and optionally bias, with optional checkpointing. + + This is the unified entry point that handles both normal execution + and checkpoint-based execution for memory efficiency. + + Args: + x_with_bias: Tuple of (x, bias) where: + - x: [s, b, C] - hidden states + - bias: [C] or None - optional bias tensor + h_post: [s, b, n] - expansion weights + manager: Optional CheckpointManager for checkpoint management. + When provided, wraps _apply_h_post with CheckpointWithoutOutput. + + Returns: + Tuple of (x_out, bias_out) where: + - x_out: [s, b, n*C] - expanded hidden states + - bias_out: [s, b, n*C] or None - expanded bias if input bias was not None + """ + x, bias = x_with_bias + + if manager is not None: + from megatron.core.tensor_parallel.random import CheckpointWithoutOutput + + # Checkpoint _apply_h_post to discard the output + x_out = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint( + self._apply_h_post, x, h_post + ) + + # Checkpoint _apply_h_post for bias if not None + if bias is not None: + bias_out = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint( + self._apply_h_post, bias, h_post + ) + else: + bias_out = None + else: + # Normal execution without checkpoint + x_out = self._apply_h_post(x, h_post) + bias_out = self._apply_h_post(bias, h_post) if bias is not None else None + + return x_out, bias_out + + def aggregate(self, x: Tensor, h_pre: Tensor) -> Tensor: + """ + Aggregate n-stream to 1-stream. + + Args: + x: [s, b, n*C] - n-stream hidden states + h_pre: [s, b, n] - aggregation weights + + Returns: + aggregated: [s, b, C] - single stream hidden states + """ + s, b, _ = x.shape + C = self.hidden_size + x_streams = x.view(s, b, self.n, C) + return self._h_aggregate_op(x_streams, h_pre) + + @torch.compile + def apply_h_res(self, h_res: Tensor, residual: Tensor) -> Tensor: + """ + Apply H_res to residual using H_res weights. + + Computes: H_res.T @ residual + + Args: + h_res: [s, b, n, n] - residual mixing matrix + residual: [s, b, n*C] - n-stream hidden states + """ + s, b, _ = residual.shape + n = self.n + C = self.hidden_size + + # Reshape for bmm: [s, b, n, n] -> [s*b, n, n] + h_res_batched = h_res.view(s * b, n, n) + # [s, b, n*C] -> [s, b, n, C] -> [s*b, n, C] + residual_batched = residual.view(s, b, n, C).view(s * b, n, C) + + # Batch matrix multiply: [s*b, n, n].T @ [s*b, n, C] -> [s*b, n, C] + mixed = torch.bmm(h_res_batched.transpose(1, 2), residual_batched) + + return mixed.view(s, b, n * C) + + def forward( + self, + hidden_states: Tensor, + mhc_recompute_manager: Optional['CheckpointManager'] = None, + return_residual: bool = False, + ) -> Tuple[Tensor, ...]: + """ + Full mHC forward pass. + + Uses BroadcastTensorFused to split hidden_states into 3 autograd-graph + children so that gradient accumulation from the 3 consumers + (compute_mappings, aggregate, fused_h_res_h_post_bda) is handled by a + single fused add instead of PyTorch's default sequential accumulation. + + Args: + hidden_states: [s, b, n*C] - n-stream hidden states + mhc_recompute_manager: Optional CheckpointManager for checkpoint management. + When provided, uses _forward_with_checkpoint for memory-efficient execution. + + Returns: + The compatible 3-tuple ``(aggregated, h_res, h_post)`` by default. + HybridModel callers set ``return_residual=True`` to also receive the + residual branch created by ``BroadcastTensorFused``. + aggregated: [s, b, C] - aggregated input for layer computation + h_res: [s, b, n, n] - residual mixing matrix (for fused kernel) + h_post: [s, b, n] - expansion weights + residual: [s, b, n*C] - residual view for fused_h_res_h_post_bda + """ + if mhc_recompute_manager is not None: + result = self._forward_with_checkpoint(hidden_states, mhc_recompute_manager) + else: + result = self._forward_normal(hidden_states) + return result if return_residual else result[:3] + + def _forward_normal(self, hidden_states: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + """ + Normal forward pass without checkpointing. + + Args: + hidden_states: [s, b, n*C] - n-stream hidden states + + Returns: + aggregated: [s, b, C] - aggregated input for layer computation + h_res: [s, b, n, n] - residual mixing matrix (for fused kernel) + h_post: [s, b, n] - expansion weights + residual: [s, b, n*C] - residual view for fused_h_res_h_post_bda + """ + # Split into 3 views to avoid extra grad accumulations in backward + hs_for_mappings, hs_for_aggregate, hs_for_residual = BroadcastTensorFused.apply( + hidden_states, self._fused_add_3_op + ) + + # Compute mappings + h_pre, h_post, h_res = self.compute_mappings(hs_for_mappings) + + # Aggregate for layer input + with torch.cuda.nvtx.range("HyperConnection::aggregate"): + aggregated = self.aggregate(hs_for_aggregate, h_pre) + + return aggregated, h_res, h_post, hs_for_residual + + def _forward_with_checkpoint( + self, hidden_states: Tensor, manager: 'CheckpointManager' + ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + """ + Forward pass with checkpointing for memory efficiency. + + compute_mappings is called directly (not checkpointed) since its outputs + (h_pre, h_post, h_res) are needed downstream. Only aggregate is wrapped with + CheckpointWithoutOutput and auto-registered to the manager. + apply_h_res is deferred to fused_h_res_h_post_bda for kernel fusion. + + Args: + hidden_states: [s, b, n*C] - n-stream hidden states + manager: CheckpointManager for unified recomputation + + Returns: + aggregated: [s, b, C] - aggregated input for layer computation + h_res: [s, b, n, n] - residual mixing matrix (for fused kernel) + h_post: [s, b, n] - expansion weights + residual: [s, b, n*C] - residual view for fused_h_res_h_post_bda + """ + from megatron.core.tensor_parallel.random import CheckpointWithoutOutput + + # Split into 3 views to avoid extra grad accumulations in backward + hs_for_mappings, hs_for_aggregate, hs_for_residual = BroadcastTensorFused.apply( + hidden_states, self._fused_add_3_op + ) + + h_pre, h_post, h_res = self.compute_mappings(hs_for_mappings) + + # Checkpoint aggregate - auto-registers to manager + aggregated = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint( + self.aggregate, hs_for_aggregate, h_pre + ) + + return aggregated, h_res, h_post, hs_for_residual + + # ==================== Block-level utilities ==================== + + @staticmethod + def input_expand(x: Tensor, n: int) -> Tensor: + """ + Expand 1-stream to n-stream at TransformerBlock entry. + + Simple replication strategy: each stream initialized as a copy of input. + + Args: + x: [s, b, C] - single stream hidden states + n: Number of residual streams + + Returns: + expanded: [s, b, n*C] - n-stream hidden states + """ + s, b, C = x.shape + # Replicate input to n streams + expanded = x.unsqueeze(2).expand(s, b, n, C).contiguous() + return expanded.view(s, b, n * C) + + @staticmethod + def output_contract(x: Tensor, n: int) -> Tensor: + """ + Contract n-stream to 1-stream at TransformerBlock exit. + + Simple averaging strategy: average all streams. + + Args: + x: [s, b, n*C] - n-stream hidden states + n: Number of residual streams + + Returns: + contracted: [s, b, C] - single stream hidden states + """ + s, b, nC = x.shape + C = nC // n + # Average all streams + x_streams = x.view(s, b, n, C) + contracted = x_streams.mean(dim=2) + return contracted + + # ==================== Fused kernel placeholder ==================== + + @nvtx_decorator(message="HyperConnection::fused_h_res_h_post_bda") + def fused_h_res_h_post_bda( + self, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + layer_output_with_bias: Tuple[Tensor, Optional[Tensor]], + dropout_prob: float, + training: bool, + fused: bool, + manager: Optional['CheckpointManager'] = None, + ) -> Tensor: + """ + Fused kernel combining apply_h_res, apply_h_post and bias-dropout-add. + + This is a placeholder for future kernel fusion optimization. + Currently implements the operations sequentially using native PyTorch. + + The computation flow is: + 1. mixed = H_res.T @ original_residual (apply_h_res) + 2. expanded = H_post^T @ layer_output (apply_h_post) + 3. output = dropout(expanded + bias) + mixed (bias-dropout-add) + + Args: + h_res: [s, b, n, n] - residual mixing matrix + original_residual: [s, b, n*C] - n-stream hidden states (before H_res applied) + h_post: [s, b, n] - expansion weights + layer_output_with_bias: Tuple of (x, bias) where: + - x: [s, b, C] - layer output (attention or MLP output) + - bias: [C] or None - optional bias tensor + dropout_prob: Dropout probability + training: Whether in training mode + fused: Whether to use fused BDA implementation + manager: Optional CheckpointManager for checkpoint management. + When provided, each operation is wrapped with CheckpointWithoutOutput. + + Returns: + output: [s, b, n*C] - final output after all operations + """ + if manager is not None: + return self._fused_h_res_h_post_bda_with_checkpoint( + h_res, + original_residual, + h_post, + layer_output_with_bias, + dropout_prob, + training, + fused, + manager, + ) + else: + return self._fused_h_res_h_post_bda_native( + h_res, + original_residual, + h_post, + layer_output_with_bias, + dropout_prob, + training, + fused, + ) + + def _fused_h_res_h_post_bda_native( + self, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + layer_output_with_bias: Tuple[Tensor, Optional[Tensor]], + dropout_prob: float, + training: bool, + fused: bool, + ) -> Tensor: + """ + h_res, h_post and bda. + + When dropout is zero (or inference), uses a single fused/reference kernel + for H_res.T @ residual + H_post * (x + bias). Falls back to unfused + implementation when dropout is needed. + + Args: + h_res: [s, b, n, n] - residual mixing matrix + original_residual: [s, b, n*C] - n-stream hidden states + h_post: [s, b, n] - expansion weights + layer_output_with_bias: Tuple of (x, bias) + dropout_prob: Dropout probability + training: Whether in training mode + fused: Whether to use fused BDA implementation + + Returns: + output: [s, b, n*C] - final output + """ + x, bias = layer_output_with_bias + + if dropout_prob == 0.0 or not training: + s, b, _ = original_residual.shape + n = self.n + C = self.hidden_size + orig_reshaped = original_residual.view(s, b, n, C) + output = self._h_post_bda_op(h_res, orig_reshaped, h_post, x, bias) + return output.view(s, b, n * C) + + from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add + + with torch.cuda.nvtx.range("HyperConnection::apply_h_res"): + mixed = self.apply_h_res(h_res, original_residual) + with torch.cuda.nvtx.range("HyperConnection::apply_h_post"): + x_expanded = self._apply_h_post(x, h_post) + bias_expanded = self._apply_h_post(bias, h_post) if bias is not None else None + bda_func = get_bias_dropout_add(training, fused) + with torch.cuda.nvtx.range("HyperConnection::bda"): + output = bda_func((x_expanded, bias_expanded), mixed, dropout_prob) + return output + + @nvtx_decorator(message="HyperConnection::fused_h_res_h_post_bda_with_checkpoint") + def _fused_h_res_h_post_bda_with_checkpoint( + self, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + layer_output_with_bias: Tuple[Tensor, Optional[Tensor]], + dropout_prob: float, + training: bool, + fused: bool, + manager: 'CheckpointManager', + ) -> Tensor: + """ + Checkpointed variant of _fused_h_res_h_post_bda_native. + + Wraps compute in CheckpointWithoutOutput for activation memory savings. + Cannot reuse _native directly because checkpoint requires all args to be + positional Tensors; tuple/Optional/scalar args are unpacked or captured + via closure instead. + + Args: + h_res: [s, b, n, n] - residual mixing matrix + original_residual: [s, b, n*C] - n-stream hidden states + h_post: [s, b, n] - expansion weights + layer_output_with_bias: Tuple of (x, bias) + dropout_prob: Dropout probability + training: Whether in training mode + fused: Whether to use fused BDA implementation + manager: CheckpointManager for checkpoint management + + Returns: + output: [s, b, n*C] - final output + """ + from megatron.core.tensor_parallel.random import CheckpointWithoutOutput + + x, bias = layer_output_with_bias + n = self.n + C = self.hidden_size + + # Fast path: no dropout — use fused/reference h_post_bda kernel (same as _native) + if dropout_prob == 0.0 or not training: + + def _fused_wrapper(h_res, original_residual, h_post, x, *optional_bias): + s, b, _ = original_residual.shape + orig_reshaped = original_residual.view(s, b, n, C) + b_arg = optional_bias[0] if optional_bias else None + return self._h_post_bda_op(h_res, orig_reshaped, h_post, x, b_arg).view(s, b, n * C) + + ckpt = CheckpointWithoutOutput(ckpt_manager=manager) + if bias is not None: + output = ckpt.checkpoint(_fused_wrapper, h_res, original_residual, h_post, x, bias) + else: + output = ckpt.checkpoint(_fused_wrapper, h_res, original_residual, h_post, x) + + # Slow path: dropout required — fused kernel does not support dropout, + # fall back to sequential apply_h_res + apply_h_post + bda + else: + from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add + + bda_func = get_bias_dropout_add(training, fused) + has_bias = bias is not None + + def _native_wrapper(h_res, original_residual, h_post, x, *optional_bias): + with torch.cuda.nvtx.range("HyperConnection::apply_h_res"): + mixed = self.apply_h_res(h_res, original_residual) + with torch.cuda.nvtx.range("HyperConnection::apply_h_post"): + x_expanded = self._apply_h_post(x, h_post) + if has_bias: + bias_expanded = self._apply_h_post(optional_bias[0], h_post) + else: + bias_expanded = None + with torch.cuda.nvtx.range("HyperConnection::bda"): + output = bda_func((x_expanded, bias_expanded), mixed, dropout_prob) + return output + + ckpt = CheckpointWithoutOutput(ckpt_manager=manager) + if has_bias: + output = ckpt.checkpoint(_native_wrapper, h_res, original_residual, h_post, x, bias) + else: + output = ckpt.checkpoint(_native_wrapper, h_res, original_residual, h_post, x) + + return output + + +# ==================== Checkpoint utilities for mHC ==================== + + +class HyperConnectionCheckpoint: + """ + Checkpoint utility for mHC intermediate activations. + + Implements the paper's "recomputing strategy" to reduce memory footprint + by discarding intermediate n-stream activations and recomputing on-the-fly. + """ + + @staticmethod + def compute_optimal_block_size(num_layers: int, num_streams: int) -> int: + """ + Compute optimal recomputation block size. + + From paper Eq. (20): L_r^* ≈ sqrt(nL/(n+2)) + + Args: + num_layers: Total number of transformer layers + num_streams: Number of residual streams (n) + + Returns: + block_size: Optimal block size for checkpointing + """ + block_size = int(math.sqrt(num_streams * num_layers / (num_streams + 2))) + return max(1, block_size) diff --git a/megatron/core/transformer/mlp.py b/megatron/core/transformer/mlp.py index 1a578151f1e..508e955d327 100644 --- a/megatron/core/transformer/mlp.py +++ b/megatron/core/transformer/mlp.py @@ -276,6 +276,7 @@ def forward( bias_parallel, per_token_scale.unsqueeze(-1), self.config.activation_func_fp8_input_store, + self.config.activation_func_clamp_value, ) elif self.activation_func == quick_gelu and self.config.gated_linear_unit: intermediate_parallel = weighted_bias_quick_geglu_impl( @@ -307,6 +308,7 @@ def forward( self.config.cpu_offloading and self.config.cpu_offloading_activations and HAVE_TE, + self.config.activation_func_clamp_value, ) else: raise ValueError("Only support fusion of gelu and swiglu") diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py index 558b1b07a15..382b9ca71b6 100644 --- a/megatron/core/transformer/module.py +++ b/megatron/core/transformer/module.py @@ -234,6 +234,13 @@ def _te_cuda_graph_backward_dw_graph(self, microbatch_idx): return self.cuda_graphs[cg_index].backward_dw() + def _is_thd_cuda_graph(self): + """Return whether this layer uses static THD Transformer Engine graphs.""" + return ( + getattr(self.config, 'sequence_packing_scheduler', None) is not None + and self.config.cuda_graph_impl == "transformer_engine" + ) + def get_layer_static_inputs(self, seq_length, micro_batch_size): """ Get the static inputs for the layer. @@ -245,22 +252,35 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): Dict[str, torch.Tensor]: A dictionary containing the static inputs for the layer. """ # Calculate data shape related values. - context_parallel_size = self.config.context_parallel_size - slen_per_cp = seq_length // context_parallel_size sequence_parallel = self.config.sequence_parallel tensor_model_parallel_size = self.config.tensor_model_parallel_size - slen_per_cptp = ( - slen_per_cp // tensor_model_parallel_size if sequence_parallel else slen_per_cp - ) - - static_inputs = {} - static_inputs["hidden_states"] = torch.ones( - (slen_per_cptp, micro_batch_size, self.config.hidden_size), - dtype=torch.bfloat16, - requires_grad=True, - device=torch.cuda.current_device(), - ) - return static_inputs + if self._is_thd_cuda_graph(): + assert self.config.max_seqlen_per_dp_cp_rank is not None, ( + "max_seqlen_per_dp_cp_rank must be set for THD CUDA graphs." + ) + sequence_length = self.config.max_seqlen_per_dp_cp_rank + batch_size = 1 + else: + sequence_length = seq_length // self.config.context_parallel_size + batch_size = micro_batch_size + if sequence_parallel: + sequence_length //= tensor_model_parallel_size + + if self.config.bf16: + dtype = torch.bfloat16 + elif self.config.fp16: + dtype = torch.float16 + else: + dtype = torch.float32 + + return { + "hidden_states": torch.ones( + (sequence_length, batch_size, self.config.hidden_size), + dtype=dtype, + requires_grad=True, + device=torch.cuda.current_device(), + ) + } def setup_manual_hooks(self, make_hook_func): """ @@ -433,6 +453,40 @@ def float_conversion(val): return conversion_helper(val, float_conversion) +def mark_keep_in_fp32(tensor: torch.Tensor) -> torch.Tensor: + """Mark a parameter or buffer so that ``Float16Module`` keeps it in FP32. + + Args: + tensor: The parameter or buffer to mark. + + Returns: + The same tensor, for call-site convenience. + """ + tensor.keep_in_fp32 = True + return tensor + + +def convert_module_to_dtype_except_fp32_marked( + module: torch.nn.Module, dtype: torch.dtype +) -> torch.nn.Module: + """Cast floating-point parameters and buffers except those marked to stay in FP32. + + Args: + module: The module to convert in place. + dtype: The target floating-point dtype. + + Returns: + The converted module. + """ + return module._apply( + lambda tensor: ( + tensor.to(dtype) + if tensor.is_floating_point() and not getattr(tensor, 'keep_in_fp32', False) + else tensor + ) + ) + + class Float16Module(MegatronModule): """Float 16 Module. @@ -455,13 +509,17 @@ def __init__(self, config: TransformerConfig, module: torch.nn.Module): self.pg_collection = getattr(module, 'pg_collection', None) if self.fp16: - self.add_module('module', module.half()) + self.add_module( + 'module', convert_module_to_dtype_except_fp32_marked(module, torch.half) + ) def float16_convertor(val): return val.half() elif self.bf16: - self.add_module('module', module.bfloat16()) + self.add_module( + 'module', convert_module_to_dtype_except_fp32_marked(module, torch.bfloat16) + ) def float16_convertor(val): return val.bfloat16() diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index d60a83b9af7..35b9b6deabf 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -368,7 +368,14 @@ def _is_fused_impl_supported(self) -> bool: if not (use_glu_fusion or use_srelu_fusion): return False if self.config.activation_func == F.silu: - pass + if self.config.activation_func_clamp_value is not None: + if not is_te_min_version("2.17.0.dev0"): + return False + try: + from transformer_engine.pytorch.ops import ScaledClampedQGeGLU # noqa: F401 + except ImportError: + return False + return True elif self.config.activation_func == quick_gelu: try: from transformer_engine.pytorch.ops import ScaledClampedQGeGLU # noqa: F401 @@ -470,20 +477,35 @@ def register_grouped_linear_params( ) ops.append(op) - # Activation and post-multiply probs (SwiGLU, clamped quick-GeGLU, or SReLU) + # Activation and post-multiply probs (SwiGLU, clamped GLU, or SReLU). glu_interleave = self.config.moe_mlp_glu_interleave_size activation_recompute_in_mlp = bool(getattr(self, "activation_recompute", False)) if self.config.activation_func == F.silu and self.config.gated_linear_unit: - if ( - "activation_recompute_in_mlp" - in inspect.signature(te.pytorch.ops.ScaledSwiGLU).parameters - ): - op = te.pytorch.ops.ScaledSwiGLU( - glu_interleave_size=glu_interleave, - activation_recompute_in_mlp=activation_recompute_in_mlp, - ) + clamp_value = self.config.activation_func_clamp_value + if clamp_value is not None: + clamped_glu_kwargs = { + "glu_interleave_size": glu_interleave, + "alpha": 1.0, + "limit": clamp_value, + "glu_linear_offset": 0.0, + } + if ( + "activation_recompute_in_mlp" + in inspect.signature(te.pytorch.ops.ScaledClampedQGeGLU).parameters + ): + clamped_glu_kwargs["activation_recompute_in_mlp"] = activation_recompute_in_mlp + op = te.pytorch.ops.ScaledClampedQGeGLU(**clamped_glu_kwargs) else: - op = te.pytorch.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave) + if ( + "activation_recompute_in_mlp" + in inspect.signature(te.pytorch.ops.ScaledSwiGLU).parameters + ): + op = te.pytorch.ops.ScaledSwiGLU( + glu_interleave_size=glu_interleave, + activation_recompute_in_mlp=activation_recompute_in_mlp, + ) + else: + op = te.pytorch.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave) elif self.config.activation_func == quick_gelu and self.config.gated_linear_unit: clamp = self.config.activation_func_clamp_value if clamp is not None: @@ -776,6 +798,7 @@ def bias_act_func(intermediate_parallel, bias_parallel, permuted_probs): bias_parallel, permuted_probs, self.config.activation_func_fp8_input_store, + self.config.activation_func_clamp_value, ) elif self.activation_func == quick_gelu and self.config.gated_linear_unit: intermediate_parallel = weighted_bias_quick_geglu_impl( diff --git a/megatron/core/transformer/moe/fused_a2a.py b/megatron/core/transformer/moe/fused_a2a.py index c281c51b4fb..0a98e05c401 100644 --- a/megatron/core/transformer/moe/fused_a2a.py +++ b/megatron/core/transformer/moe/fused_a2a.py @@ -52,9 +52,12 @@ def get_buffer(group: torch.distributed.ProcessGroup, hidden_bytes: int): num_nvl_bytes = max( config.get_nvl_buffer_size_hint(hidden_bytes, group.size()), num_nvl_bytes ) - num_rdma_bytes = max( - config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), num_rdma_bytes - ) + # Local-only EP groups do not need an RDMA buffer, and DeepEP builds + # without internode support may not expose RDMA size hints. + if group.size() > torch.cuda.device_count(): + num_rdma_bytes = max( + config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), num_rdma_bytes + ) # Allocate buffer if not existed or not enough buffer # NOTES: the adaptive routing configuration of the network **must be off** @@ -277,6 +280,10 @@ def set_deepep_num_sms(num_sms): _hybrid_ep_buffer = None +# HybridEP dispatch/combine kernels use 64-token chunks for their public APIs. +HYBRIDEP_TOKEN_ALIGNMENT = 64 + + def init_hybrid_ep_buffer( group: torch.distributed.ProcessGroup, hidden_dim: int, diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index deebd3472ea..db6f46d2dca 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -124,7 +124,12 @@ def __call__( class RouterInterface(Protocol): """Interface for the router used in an MoELayer.""" - def forward(self, input: torch.Tensor, /) -> tuple[torch.Tensor, torch.Tensor]: + def forward( + self, + input: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + input_ids: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: """Forward pass of the router. Returns: @@ -144,7 +149,12 @@ class RouterBuilder(Protocol): """Protocol for building a Router.""" def __call__( - self, /, *, config: TransformerConfig, pg_collection: ProcessGroupCollection | None + self, + /, + *, + config: TransformerConfig, + pg_collection: ProcessGroupCollection | None, + is_mtp_layer: bool = False, ) -> RouterInterface: ... @@ -259,6 +269,8 @@ def __init__( self.router = self.submodules.router( config=self.config, pg_collection=pg_collection, is_mtp_layer=is_mtp_layer ) + if layer_number is not None: + self.router.set_layer_number(layer_number) self.tp_group = pg_collection.tp # Initialize latent projections. @@ -435,13 +447,23 @@ def setup_delayed_wgrad_for_dispatch_backward_overlap(self): self._delayed_wgrad_stream = torch.cuda.Stream(device="cuda") @maybe_skip_or_early_return_by_cudagraph("route") - def route(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tensor] = None): + def route( + self, + hidden_states: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + input_ids: Optional[torch.Tensor] = None, + ): """Compute token routing for preprocessing. This method uses the router to determine which experts to send each token to, producing routing probabilities and a mapping. """ - probs, routing_map = apply_module(self.router)(hidden_states, padding_mask) + if input_ids is None: + probs, routing_map = apply_module(self.router)(hidden_states, padding_mask) + else: + probs, routing_map = apply_module(self.router)( + hidden_states, padding_mask, input_ids=input_ids + ) return probs, routing_map @maybe_skip_or_early_return_by_cudagraph("preprocess") @@ -600,6 +622,7 @@ def forward( hidden_states: torch.Tensor, intermediate_tensors=None, padding_mask: Optional[torch.Tensor] = None, + input_ids: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: """Forward pass for the MoE layer. @@ -611,9 +634,11 @@ def forward( Args: hidden_states (torch.Tensor): The input tensor shape [seq_length, bsz, hidden_size]. - padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens. - Shape [seq_length, bsz]. True for valid tokens, - False for padding tokens. Defaults to None. + padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions. + Shape [bsz, seq_length]. True for padding, + False for valid tokens. Defaults to None. + input_ids (torch.Tensor, optional): Token IDs with shape + [batch_size, seq_length]. Required by hash routing. Returns: A tuple containing the output tensor and the MLP bias, if any. """ @@ -634,6 +659,28 @@ def forward( else: self.token_dispatcher = self._training_token_dispatcher self.shared_expert_overlap = self.config.moe_shared_expert_overlap + + # Pipeline stages receive a CP-local mask, while sequence-parallel + # activations are additionally sharded across TP ranks. + if padding_mask is not None and padding_mask.shape[1] != hidden_states.shape[0]: + if ( + self.config.sequence_parallel + and padding_mask.shape[1] % self.attn_tp_group.size() == 0 + and padding_mask.shape[1] // self.attn_tp_group.size() + == hidden_states.shape[0] + ): + padding_mask = ( + tensor_parallel.scatter_to_sequence_parallel_region( + padding_mask.transpose(0, 1).contiguous(), group=self.attn_tp_group + ) + .transpose(0, 1) + .contiguous() + ) + else: + raise AssertionError( + f"padding_mask shape {padding_mask.shape} cannot be aligned to " + f"hidden_states sequence length {hidden_states.shape[0]}" + ) # Transpose from [bsz, seq_length] to [seq_length, bsz] to align with hidden_states if padding_mask is not None: padding_mask = padding_mask.transpose(0, 1).bool() @@ -643,7 +690,7 @@ def custom_forward(hidden_states, intermediate_tensors=None, padding_mask=None): try: if "route" in self.fwd_execution_map: shared_expert_output = self.shared_experts_compute(hidden_states) - probs, routing_map = self.route(hidden_states, padding_mask) + probs, routing_map = self.route(hidden_states, padding_mask, input_ids) hidden_states, probs = self.preprocess(hidden_states, probs, routing_map) if intermediate_tensors is not None: diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index e4591ce3acf..0e60882a92b 100644 --- a/megatron/core/transformer/moe/router.py +++ b/megatron/core/transformer/moe/router.py @@ -1,7 +1,7 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from abc import ABC, abstractmethod -from typing import Optional, Union +from typing import Optional, Sequence, Union import torch @@ -22,7 +22,6 @@ sinkhorn, switch_load_balancing_loss_func, topk_routing_with_score_function, - z_loss_func, ) from megatron.core.transformer.moe.router_replay import RouterReplay from megatron.core.transformer.transformer_config import TransformerConfig @@ -173,6 +172,9 @@ def __init__( self.score_function = self.config.moe_router_score_function self.input_jitter = None self.frozen_expert_bias = False + self.mtp_layer_number: Optional[int] = None + self.is_hash_layer = False + self.register_buffer('tid2eid', None) self.enable_expert_bias = self.config.moe_router_enable_expert_bias if self.enable_expert_bias: @@ -256,6 +258,29 @@ def __init__( if self.config.moe_enable_routing_replay: self.router_replay = RouterReplay() + def set_layer_number(self, layer_number: int): + """Set the layer number and initialize hash routing for eligible layers.""" + super().set_layer_number(layer_number) + self.is_hash_layer = ( + not self.is_mtp_layer + and self.config.moe_n_hash_layers > 0 + and layer_number <= self.config.moe_n_hash_layers + ) + if not self.is_hash_layer: + return + + if self.tid2eid is None: + token_ids = torch.arange(self.config.actual_vocab_size, device=self.weight.device) + expert_offsets = torch.arange(self.topk, device=token_ids.device) + self.tid2eid = ((token_ids[:, None] + expert_offsets) % self.num_experts).to( + torch.int32 + ) + + # Dynamic expert bias applies to learned top-k selection, not a fixed lookup table. + self.enable_expert_bias = False + self.local_tokens_per_expert = None + self.expert_bias = None + def _maintain_float32_expert_bias(self): """ Maintain the expert bias in float32. @@ -441,6 +466,8 @@ def _apply_aux_loss( "load_balancing_loss", self.tp_cp_group, valid_token_count=local_num_tokens, + aux_loss_scale_reduce_groups=(self.tp_cp_group,), + aux_loss_scale_num_tokens=total_num_tokens, ) return probs @@ -498,6 +525,8 @@ def _apply_seq_aux_loss( # local_num_tokens is per-sequence (bsz folded into the expert dim above); # * bsz recovers the micro-batch total, else per-token-loss scaling keeps a 1/MBS. valid_token_count=local_num_tokens * bsz, + aux_loss_scale_reduce_groups=(self.tp_cp_group,), + aux_loss_scale_num_tokens=total_num_tokens * bsz, ) return probs @@ -544,6 +573,11 @@ def _apply_global_aux_loss( self.tp_dp_cp_group, needs_dp_avg=False, valid_token_count=local_num_tokens, + # The global aux-loss statistics/logging domain is TP x DP x CP, but + # per-token-loss gradient normalization already reduces the denominator + # across DP x CP in finalize_model_grads. Scale the aux-loss numerator + # over TP x CP only. + aux_loss_scale_reduce_groups=(self.tp_cp_group,), ) return probs @@ -553,9 +587,13 @@ def attach_and_log_load_balancing_loss( aux_loss_coeff: float, aux_loss: torch.Tensor, aux_loss_name: str, - reduce_group: torch.distributed.ProcessGroup, + reduce_group: Optional[torch.distributed.ProcessGroup], + avg_group: Optional[torch.distributed.ProcessGroup] = None, needs_dp_avg: bool = True, valid_token_count: Optional[Union[int, torch.Tensor]] = None, + aux_loss_logging_reduce_groups: Optional[Sequence[torch.distributed.ProcessGroup]] = None, + aux_loss_scale_reduce_groups: Optional[Sequence[torch.distributed.ProcessGroup]] = None, + aux_loss_scale_num_tokens: Optional[Union[int, torch.Tensor]] = None, ): """Attach aux loss function to activation and add to logging. @@ -564,11 +602,20 @@ def attach_and_log_load_balancing_loss( aux_loss_coeff (float): Coefficient for the aux loss. aux_loss (torch.Tensor): Computed aux loss. aux_loss_name (str): Name of the aux loss for logging. - reduce_group (torch.distributed.ProcessGroup): Process group for reduction. + reduce_group (torch.distributed.ProcessGroup, optional): Process group for deferred + logging reduction. + avg_group (torch.distributed.ProcessGroup, optional): Process group for deferred + logging average. needs_dp_avg (bool): Whether to average this metric across DP ranks after reduce_group. valid_token_count (int or torch.Tensor, optional): Number of valid tokens excluding padding tokens. Can be a Python int or a torch.Tensor (typically 0-d tensor). If None, uses activation.shape[0]. Defaults to None. + aux_loss_logging_reduce_groups (Sequence[torch.distributed.ProcessGroup], optional): + Process groups to reduce the detached logging value over before recording it. + aux_loss_scale_reduce_groups (Sequence[torch.distributed.ProcessGroup], optional): + Process groups to reduce the local valid-token count over for per-token scaling. + aux_loss_scale_num_tokens (int or torch.Tensor, optional): Pre-reduced valid-token + count for per-token scaling. """ # When using repeated MTP layers, the loss is counted "mtp_num_layers" times. # To avoid accumulating the load balancing loss multiple times, we scale it by @@ -589,48 +636,53 @@ def attach_and_log_load_balancing_loss( num_layers += self.config.mtp_num_layers if self.is_mtp_layer: - layer_number = self.layer_number + self.config.num_layers + # Hybrid MTP depths can contain multiple internal sublayers (for example `/WE`). + # Metrics are allocated per MTP depth, not per internal hybrid sublayer. + mtp_layer_number = self.mtp_layer_number or self.layer_number + if self.config.mtp_num_layers is not None: + mtp_layer_number = min(mtp_layer_number, self.config.mtp_num_layers) + layer_number = mtp_layer_number + self.config.num_layers else: layer_number = self.layer_number + metric_value = aux_loss / aux_loss_coeff + if aux_loss_logging_reduce_groups is not None: + metric_value = metric_value.detach().clone() + for group in aux_loss_logging_reduce_groups: + torch.distributed.all_reduce(metric_value, group=group) + get_moe_metrics_tracker().record( aux_loss_name, - aux_loss / aux_loss_coeff, + metric_value, layer_number, num_layers, reduce_group=reduce_group, + avg_group=avg_group, needs_dp_avg=needs_dp_avg, ) if self.calculate_per_token_loss: - # Target final scaling on aux_loss gradients: 1 / (num_micro_batches * dp_size), - # matching the !calculate_per_token_loss path. - # - # --calculate-per-token-loss already divides every parameter gradient by - # total_global_tokens (the global non-padded token count summed in - # finalize_model_grads). The router's `num_local_tokens` (= activation.shape[0]) - # is sequence-parallel sharded — the router weight is marked - # `sequence_parallel=True` in Router.reset_parameters (see - # `setattr(self.weight, 'sequence_parallel', ...)` above), so each TP rank - # computes a partial gradient on the router weight from its local sequence - # shard, and `_allreduce_non_tensor_model_parallel_grads` SUMS those partial - # gradients across the TP group. Re-expressing total_global_tokens in terms of the - # router's `num_local_tokens`: - # total_global_tokens - # = num_micro_batches * dp_cp_size * loss_func_local_tokens - # = num_micro_batches * dp_cp_size * tp_size * num_local_tokens - # = num_micro_batches * dp_size * (num_local_tokens * tp_cp_group.size()) - # (using loss_func_local_tokens = tp_size * num_local_tokens, then regrouping - # dp_cp_size * tp_size as dp_size * tp_cp_group.size()). - # - # So pre-multiplying aux_loss by num_local_tokens * tp_cp_group.size() cancels - # that same factor in total_global_tokens above, leaving 1 / (num_micro_batches * - # dp_size) as the effective scaling on the aux_loss gradient — the target. - # Use valid_token_count (excluding padding) if provided, otherwise use total tokens. - num_local_tokens = ( - valid_token_count if valid_token_count is not None else activation.shape[0] - ) + # finalize_model_grads divides by the global non-padded token count. Multiply + # by the exact valid-token count from this aux-loss domain. With THD padding, + # local counts can differ by rank, so local_count * group_size is not sufficient. + if aux_loss_scale_num_tokens is None: + num_local_tokens = ( + valid_token_count if valid_token_count is not None else activation.shape[0] + ) + if torch.is_tensor(num_local_tokens): + aux_loss_scale_num_tokens = num_local_tokens.clone().to( + device=activation.device + ) + else: + aux_loss_scale_num_tokens = torch.tensor( + num_local_tokens, device=activation.device + ) + if aux_loss_scale_reduce_groups is None: + assert reduce_group is not None, "reduce_group is required for aux-loss scaling" + aux_loss_scale_reduce_groups = (reduce_group,) + for group in aux_loss_scale_reduce_groups: + torch.distributed.all_reduce(aux_loss_scale_num_tokens, group=group) activation = MoEAuxLossAutoScaler.apply( - activation, aux_loss * num_local_tokens * self.tp_cp_group.size() + activation, aux_loss * aux_loss_scale_num_tokens ) else: activation = MoEAuxLossAutoScaler.apply(activation, aux_loss) @@ -642,61 +694,60 @@ def apply_z_loss(self, logits, padding_mask: Optional[torch.Tensor] = None): Args: logits (torch.Tensor): The logits of the router. - padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens. - Shape in [num_tokens]. True for valid tokens, - False for padding tokens. Defaults to None. + padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions. + Shape [num_tokens]. True = padding, + False = valid. Defaults to None. Returns: torch.Tensor: The logits after applying the z-loss. """ if self.config.moe_z_loss_coeff is not None and self.training and torch.is_grad_enabled(): # Skip Z loss calculations when using torch.no_grad() or checkpointing. - moe_z_loss_coeff = self.config.moe_z_loss_coeff / self.tp_cp_group.size() - z_loss = z_loss_func(logits, moe_z_loss_coeff, padding_mask=padding_mask) - if self.calculate_per_token_loss: - # Same derivation as in attach_and_log_load_balancing_loss: - # - Target final scaling on z_loss gradients: 1 / (num_micro_batches * dp_size). - # - In terms of the router's `num_local_tokens`, the total_global_tokens - # divisor that finalize_model_grads applies factors as - # num_micro_batches * dp_size * (num_local_tokens * tp_cp_group.size()). - # - Pre-multiplying z_loss by num_local_tokens * tp_cp_group.size() cancels - # that same factor in total_global_tokens, leaving - # 1 / (num_micro_batches * dp_size) as the effective scaling — the target. - # The /tp_cp_group.size() on moe_z_loss_coeff above is a separate forward-side - # correction: z_loss is computed independently on each TP+CP rank's local - # logits and must be averaged across TP+CP rather than summed. - # Count valid tokens: sum of inverted mask (False -> True = valid) - num_local_tokens = ( - (~padding_mask).sum() if padding_mask is not None else logits.shape[0] - ) - logits = MoEAuxLossAutoScaler.apply( - logits, z_loss * num_local_tokens * self.tp_cp_group.size() - ) + logsum = torch.logsumexp(logits, dim=-1) + z_loss_values = torch.square(logsum) + if padding_mask is not None: + valid_mask = ~padding_mask + z_loss_values = z_loss_values * valid_mask + num_local_tokens = valid_mask.sum() else: - logits = MoEAuxLossAutoScaler.apply(logits, z_loss) + num_local_tokens = torch.tensor(logits.shape[0], device=logits.device) + + z_loss_sum = z_loss_values.sum() + z_loss_mean = z_loss_sum / torch.clamp(num_local_tokens, min=1) - # When using repeated MTP layers, the same MTP layer is called mtp_num_layers times. - # To avoid accumulating the z_loss multiple times, we scale it by 1/mtp_num_layers - # so the total loss is correct. + mtp_loss_scale = 1 if ( self.is_mtp_layer and self.config.mtp_use_repeated_layer and self.config.mtp_num_layers is not None ): - z_loss = z_loss / self.config.mtp_num_layers + mtp_loss_scale = self.config.mtp_num_layers + + if self.calculate_per_token_loss: + # The global per-token denominator is applied later; attach the local + # numerator so valid tokens, including rank-skewed THD padding, weight exactly. + z_loss = z_loss_sum * self.config.moe_z_loss_coeff / mtp_loss_scale + logits = MoEAuxLossAutoScaler.apply(logits, z_loss) + else: + moe_z_loss_coeff = self.config.moe_z_loss_coeff / self.tp_cp_group.size() + z_loss = z_loss_mean * moe_z_loss_coeff / mtp_loss_scale + logits = MoEAuxLossAutoScaler.apply(logits, z_loss) num_layers = self.config.num_layers if self.config.mtp_num_layers is not None: num_layers += self.config.mtp_num_layers if self.is_mtp_layer: - layer_number = self.layer_number + self.config.num_layers + mtp_layer_number = self.mtp_layer_number or self.layer_number + if self.config.mtp_num_layers is not None: + mtp_layer_number = min(mtp_layer_number, self.config.mtp_num_layers) + layer_number = mtp_layer_number + self.config.num_layers else: layer_number = self.layer_number get_moe_metrics_tracker().record( "z_loss", - z_loss / moe_z_loss_coeff, + z_loss_mean / mtp_loss_scale, layer_number, num_layers, avg_group=self.tp_dp_cp_group, @@ -735,17 +786,69 @@ def _apply_expert_bias( if self.enable_expert_bias and torch.is_grad_enabled(): with torch.no_grad(): if padding_mask is not None: - routing_map = routing_map & (~padding_mask) + flat_mask = padding_mask.reshape(-1) + assert ( + flat_mask.shape[0] == routing_map.shape[0] + ), f"padding_mask flat {flat_mask.shape} vs routing_map {routing_map.shape}" + routing_map = routing_map & (~flat_mask).unsqueeze(-1) self.local_tokens_per_expert += routing_map.sum(dim=0) - def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = None): + def _hash_routing(self, logits: torch.Tensor, input_ids: torch.Tensor): + """Route tokens through the token-to-expert lookup table. + + Gating logits still provide the combination weights, while expert selection + comes from ``tid2eid``. + """ + if self.score_function == "softmax": + scores = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits) + elif self.score_function == "sigmoid": + scores = torch.sigmoid(logits.float()).type_as(logits) + elif self.score_function == "sqrtsoftplus": + scores = torch.nn.functional.softplus(logits.float()).sqrt().type_as(logits) + else: + raise ValueError(f"Invalid score_function: {self.score_function}") + + # Hidden states are flattened from [sequence, batch, hidden], whereas + # model token IDs arrive as [batch, sequence]. + flat_ids = input_ids.T.reshape(-1) + assert flat_ids.numel() == logits.shape[0], ( + f"input_ids contains {flat_ids.numel()} tokens, but router logits contain " + f"{logits.shape[0]}." + ) + top_indices = self.tid2eid[flat_ids].long() + if ( + self.config.moe_router_force_load_balancing + or self.config.moe_router_force_biased is not None + ): + # Benchmark forcing must override the fixed table, just as it overrides + # learned top-k routing. + top_indices = torch.topk(logits, k=self.topk, dim=1).indices + + probs = scores.gather(1, top_indices) + if self.score_function != "softmax": + probs = probs / (probs.sum(dim=-1, keepdim=True) + 1e-20) + if self.config.moe_router_topk_scaling_factor: + probs = probs * self.config.moe_router_topk_scaling_factor + + routing_probs = torch.zeros_like(logits).scatter(1, top_indices, probs) + routing_map = torch.zeros_like(logits, dtype=torch.bool).scatter(1, top_indices, True) + return routing_probs, routing_map + + def routing( + self, + logits: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + input_ids: Optional[torch.Tensor] = None, + ): """Top-k routing function Args: logits (torch.Tensor): Logits tensor after gating. - padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens. - Shape [seq_length, bsz]. True for valid tokens, - False for padding tokens. Defaults to None. + padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions. + Shape [seq_length, bsz]. True = padding, + False = valid. Defaults to None. + input_ids (torch.Tensor, optional): Token IDs with shape [batch, sequence]. + Required when this is a hash-routing layer. Returns: probs (torch.Tensor): The probabilities of token to experts assignment. @@ -763,7 +866,18 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N logits = self.apply_z_loss(logits, padding_mask=padding_mask) # Calculate probs and routing_map for token dispatching - if self.routing_type == "sinkhorn": + if self.config.moe_n_hash_layers > 0: + assert self.layer_number is not None, ( + "Hash routing requires a layer number. Construct the router through MoELayer " + "or call set_layer_number() before routing." + ) + if self.is_hash_layer: + assert input_ids is not None, ( + "input_ids is required for hash-based routing. Pass token IDs through " + "the model, transformer block, and transformer layer." + ) + probs, routing_map = self._hash_routing(logits, input_ids) + elif self.routing_type == "sinkhorn": probs, routing_map = self.sinkhorn_load_balancing(logits) elif self.routing_type == "quantile_balancing": assert ( @@ -784,6 +898,20 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N router_replay=self.router_replay, ) + # Dropless HybridEP consumes the sparse routing map directly, so exclude padding + # rows before dispatch. Other dispatchers retain their existing fixed-route + # assumptions until they support sparse routing maps end to end. + use_dropless_hybridep = ( + self.config.moe_token_dispatcher_type == "flex" + and self.config.moe_flex_dispatcher_backend == "hybridep" + and self.config.moe_expert_capacity_factor is None + and self.config.moe_expert_rank_capacity_factor is None + ) + if padding_mask is not None and use_dropless_hybridep: + valid_tokens = (~padding_mask).unsqueeze(-1) + probs = probs * valid_tokens + routing_map = routing_map & valid_tokens + # Apply token dropping to probs and routing_map. if self.config.moe_expert_capacity_factor is not None: probs, routing_map = apply_router_token_dropping( @@ -837,15 +965,22 @@ def reset_global_aux_loss_tracker(self): self.global_tokens_per_expert.zero_() self.ga_steps.zero_() - def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = None): + def forward( + self, + input: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + input_ids: Optional[torch.Tensor] = None, + ): """ Forward pass of the router. Args: input (torch.Tensor): Input tensor. - padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens. - Shape [seq_length, bsz]. True for valid tokens, - False for padding tokens. Defaults to None. + padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions. + Shape [seq_length, bsz]. True = padding, + False = valid. Defaults to None. + input_ids (torch.Tensor, optional): Token IDs with shape [batch, sequence]. + Required when this is a hash-routing layer. """ self._maintain_float32_expert_bias() @@ -863,7 +998,7 @@ def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = No logits, self.config.moe_router_force_biased, self.layer_number ) - probs, routing_map = self.routing(logits, padding_mask=padding_mask) + probs, routing_map = self.routing(logits, padding_mask=padding_mask, input_ids=input_ids) return probs, routing_map @@ -907,12 +1042,16 @@ def __init__( f"InferenceTopKRouter requires moe_router_num_groups=None, " f"got {config.moe_router_num_groups}" ) - assert config.moe_router_score_function in ["sigmoid", "softmax"], ( - f"InferenceTopKRouter requires moe_router_score_function in " - f"['sigmoid', 'softmax'], got '{config.moe_router_score_function}'" + supported_compiled_scores = ["sigmoid", "softmax"] + assert config.moe_router_score_function in supported_compiled_scores or ( + config.moe_n_hash_layers > 0 and config.moe_router_score_function == "sqrtsoftplus" + ), ( + "InferenceTopKRouter requires moe_router_score_function to be sigmoid/softmax, " + "or sqrtsoftplus for hash routing; got " + f"{config.moe_router_score_function!r}" ) - super().__init__(config=config, pg_collection=pg_collection) + super().__init__(config=config, pg_collection=pg_collection, is_mtp_layer=is_mtp_layer) @staticmethod @torch.compile @@ -969,12 +1108,18 @@ def _forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = N ) return probs.squeeze(1), top_indices.squeeze(1) - def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = None): + def forward( + self, + input: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + input_ids: Optional[torch.Tensor] = None, + ): """Simplified forward pass for inference - returns dense tensors only. Args: input (torch.Tensor): Input tensor of shape [seq_length, bsz, hidden_size]. padding_mask (torch.Tensor, optional): Not used in inference. + input_ids (torch.Tensor, optional): Token IDs used by hash routing. Returns: Tuple[torch.Tensor, torch.Tensor]: @@ -982,7 +1127,11 @@ def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = No - top_indices: Selected expert indices [num_tokens, topk] """ - if not InferenceMode.is_active(): - return super().forward(input, padding_mask) + if ( + not InferenceMode.is_active() + or self.is_hash_layer + or self.score_function == "sqrtsoftplus" + ): + return super().forward(input, padding_mask, input_ids) return self._forward(input, padding_mask) diff --git a/megatron/core/transformer/moe/shared_experts.py b/megatron/core/transformer/moe/shared_experts.py index 0f5108995cc..66b9583226a 100644 --- a/megatron/core/transformer/moe/shared_experts.py +++ b/megatron/core/transformer/moe/shared_experts.py @@ -282,6 +282,7 @@ def linear_fc1_forward_and_act(self, overlapped_comm_output=None): intermediate_parallel, bias_parallel, self.config.activation_func_fp8_input_store, + clamp_value=self.config.activation_func_clamp_value, ) else: raise ValueError("Only support fusion of gelu and swiglu") @@ -291,8 +292,13 @@ def linear_fc1_forward_and_act(self, overlapped_comm_output=None): if self.config.gated_linear_unit: def glu(x): - x = torch.chunk(x, 2, dim=-1) - return self.config.activation_func(x[0]) * x[1] + x_glu, x_linear = torch.chunk(x, 2, dim=-1) + if (clamp_value := self.config.activation_func_clamp_value) is not None: + x_glu = x_glu.clamp(min=None, max=clamp_value) + x_linear = x_linear.clamp(min=-clamp_value, max=clamp_value) + return self.config.activation_func(x_glu) * ( + x_linear + self.config.glu_linear_offset + ) intermediate_parallel = glu(intermediate_parallel) else: @@ -409,6 +415,15 @@ def _validate_fused_grouped_swiglu(self) -> None: f"fused kernel, but got activation_func={self.config.activation_func}, " f"gated_linear_unit={self.config.gated_linear_unit}." ) + if self.config.activation_func_clamp_value is not None and ( + not is_te_min_version("2.17.0.dev0") + or not hasattr(te.pytorch.ops, "ScaledClampedQGeGLU") + ): + raise RuntimeError( + f"{self.__class__.__name__} requires Transformer Engine >= 2.17.0.dev0 " + "with pytorch.ops.ScaledClampedQGeGLU when " + "activation_func_clamp_value is set." + ) if self.config.moe_shared_expert_glu_interleave_size is None: raise ValueError( f"{self.__class__.__name__} requires " @@ -445,7 +460,7 @@ def _get_fused_grouped_swiglu_recipe(self): return self._fused_grouped_swiglu_recipe def _make_fused_grouped_swiglu_ops(self) -> torch.nn.Module: - """Construct GroupedLinear(num_groups=1) -> ScaledSwiGLU -> GroupedLinear.""" + """Construct the grouped-linear shared-expert MLP operations.""" ops = te.pytorch.ops.Sequential() tp_world_size = get_pg_size(self.tp_group) rng_state_tracker_function = None @@ -468,7 +483,17 @@ def _make_fused_grouped_swiglu_ops(self) -> torch.nn.Module: op._glu_interleave_size = glu_interleave_size ops.append(op) - ops.append(te.pytorch.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size)) + clamp_value = self.config.activation_func_clamp_value + if clamp_value is None: + activation_op = te.pytorch.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size) + else: + activation_op = te.pytorch.ops.ScaledClampedQGeGLU( + glu_interleave_size=glu_interleave_size, + alpha=1.0, + limit=clamp_value, + glu_linear_offset=0.0, + ) + ops.append(activation_op) fc2_weight = self.linear_fc2.weight op = te.pytorch.ops.GroupedLinear( diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index 4c4b65679c3..c28fc99d9ca 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -19,6 +19,7 @@ ) from megatron.core.transformer.enums import CudaGraphModule from megatron.core.transformer.moe.fused_a2a import ( + HYBRIDEP_TOKEN_ALIGNMENT, ensure_nccl_ep_bootstrapped, fused_combine, fused_dispatch, @@ -1039,11 +1040,49 @@ def __init__( self.moe_expert_rank_capacity_factor = self.config.moe_expert_rank_capacity_factor self.over_budget = torch.zeros(1, dtype=torch.bool, device='cuda') + # HybridEP dispatch expects equal per-rank input sizes. When requested, + # variable token counts are padded to the group-wide max and trimmed in combine. + self._original_num_tokens: Optional[int] = None + self._padded_num_tokens: Optional[int] = None def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor): num_tokens = routing_map.shape[0] - self.routing_map = routing_map.reshape(num_tokens, self.num_experts) - self.token_probs = probs.reshape(num_tokens, self.num_experts) + self._original_num_tokens = num_tokens + + padded_num_tokens = num_tokens + if self.config.moe_hybridep_pad_uneven_dispatch_inputs: + if ( + self.config.sequence_packing_scheduler is not None + and (torch.cuda.is_current_stream_capturing() or torch.compiler.is_compiling()) + ): + # CUDA graph path: upstream sequence packing has already padded + # routing_map to a static per-rank length. Skip the collective + # and host synchronization while tracing or capturing. + padded_num_tokens = num_tokens + else: + # Use the actual tp_ep max so all ranks in the MoE communication + # group pass the same token count to HybridEP. + max_num_tokens_across_ep = torch.tensor( + [num_tokens], device=routing_map.device, dtype=torch.long + ) + torch.distributed.all_reduce( + max_num_tokens_across_ep, op=torch.distributed.ReduceOp.MAX, group=self.group + ) + padded_num_tokens = int(max_num_tokens_across_ep.item()) + padded_num_tokens += -padded_num_tokens % HYBRIDEP_TOKEN_ALIGNMENT + self._padded_num_tokens = padded_num_tokens + + routing_map = routing_map.reshape(num_tokens, self.num_experts) + probs = probs.reshape(num_tokens, self.num_experts) + if padded_num_tokens > num_tokens: + pad_rows = padded_num_tokens - num_tokens + routing_map = torch.cat( + [routing_map, routing_map.new_zeros((pad_rows, self.num_experts))], dim=0 + ) + probs = torch.cat([probs, probs.new_zeros((pad_rows, self.num_experts))], dim=0) + + self.routing_map = routing_map + self.token_probs = probs if self.moe_expert_rank_capacity_factor is not None: pad_multiple = get_align_size_for_quantization(self.config) @@ -1051,7 +1090,7 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor): # budget). Tokens above this budget are dropped inside HybridEP; dispatch then # sets overflow_flag on the handle (accumulated in over_budget in dispatch()). budget = int( - routing_map.shape[0] + padded_num_tokens * self.config.moe_router_topk * self.moe_expert_rank_capacity_factor ) @@ -1062,7 +1101,7 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor): # in dispatch) and does not drop tokens or report overflow. # Compute the capacity for each expert at the drop_and_pad mode if self.drop_and_pad: - num_out_tokens = num_tokens * self.config.moe_router_topk + num_out_tokens = padded_num_tokens * self.config.moe_router_topk # Drop and pad the input to capacity. self.capacity = get_capacity( num_tokens=num_out_tokens, @@ -1091,6 +1130,11 @@ def dispatch( self.token_probs = self.token_probs.float() # downcast or upcast if self.config.fp8 or self.config.fp4: self.pad_multiple = get_align_size_for_quantization(self.config) + if self._padded_num_tokens is not None and hidden_states.shape[0] < self._padded_num_tokens: + pad_rows = self._padded_num_tokens - hidden_states.shape[0] + hidden_states = torch.cat( + [hidden_states, hidden_states.new_zeros((pad_rows, hidden_states.shape[-1]))], dim=0 + ) dispatched_hidden, self.dispatched_probs, _, tokens_per_expert, self.handle = ( hybrid_ep_dispatch( x=hidden_states, @@ -1137,12 +1181,20 @@ def combine( pad_multiple=self.pad_multiple, fused=self.config.moe_permute_fusion_into_hybridep, ) + if ( + self._padded_num_tokens is not None + and self._original_num_tokens is not None + and hidden_states.shape[0] > self._original_num_tokens + ): + hidden_states = hidden_states[: self._original_num_tokens] # Release the used handle/num_permuted_tokens which could change in each iteration. # For drop_and_pad mode, we don't need to reset the num_permuted_tokens and # num_dispatched_tokens, because their values never change. self.handle = None if not self.drop_and_pad: self.num_permuted_tokens = None + self._original_num_tokens = None + self._padded_num_tokens = None return hidden_states def get_permuted_hidden_states_by_experts(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index 202034986db..be8b7bf9a7a 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -143,6 +143,7 @@ def __init__( cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, name: str | None = None, ) -> None: # TODO(nschank): Restructure so that the Attention initializer knows which specific @@ -155,6 +156,7 @@ def __init__( attn_mask_type=attn_mask_type, pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, name=name, ) self.config: MLATransformerConfig @@ -482,6 +484,7 @@ def __init__( cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, name: str | None = None, ): if pg_collection is None: @@ -496,6 +499,7 @@ def __init__( cp_comm_type=cp_comm_type, pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, name=name, ) @@ -716,8 +720,11 @@ def get_query_key_value_tensors( cu_seqlens_kv = packed_seq_params.cu_seqlens_kv_padded else: cu_seqlens_kv = packed_seq_params.cu_seqlens_kv + rope_max_seqlen_q = packed_seq_params.max_seqlen_q + rope_max_seqlen_kv = packed_seq_params.max_seqlen_kv else: cu_seqlens_q = cu_seqlens_kv = None + rope_max_seqlen_q = rope_max_seqlen_kv = None # ========================================= # QKV down projection and layernorm @@ -929,6 +936,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po mscale=mscale, cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, + max_seqlen=rope_max_seqlen_q, ) # k_pos_emb:[num_tokens, 1, qk_pos_emb_head_dim] k_pos_emb = apply_rotary_pos_emb( @@ -939,6 +947,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po mscale=mscale, cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, + max_seqlen=rope_max_seqlen_kv, ) # query: [num_tokens, n, (qk_head_dim + v_head_dim)] @@ -1227,6 +1236,7 @@ def __init__( attn_mask_type=AttnMaskType.padding, cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, + is_mtp_layer: bool = False, pp_layer_offset: Optional[int] = None, name: str | None = None, ): @@ -1242,6 +1252,7 @@ def __init__( attention_type="self", cp_comm_type=cp_comm_type, pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, pp_layer_offset=pp_layer_offset, name=name, ) diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py index b20514ce6a4..81a2e0df15e 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Callable, List, Optional, Union import torch +import torch.nn as nn from torch import Tensor from megatron.core import InferenceParams, parallel_state, tensor_parallel @@ -28,7 +29,8 @@ inference_all_gather_from_tensor_model_parallel_region, ) from megatron.core.transformer.enums import AttnMaskType, LayerType -from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.hyper_connection import learned_output_contract +from megatron.core.transformer.module import MegatronModule, mark_keep_in_fp32 from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.torch_norm import LayerNormBuilder from megatron.core.transformer.transformer_block import TransformerBlockSubmodules @@ -134,7 +136,9 @@ def tie_output_layer_state_dict( ) -def roll_tensor(tensor, shifts=-1, dims=-1, cp_group=None, packed_seq_params=None): +def roll_tensor( + tensor, shifts=-1, dims=-1, cp_group=None, packed_seq_params=None, fill_value=0 +): """Roll the tensor input along the sequence dimension with Context Parallelism (CP) support. This function extends the original roll_tensor to support Context Parallelism, which allows @@ -157,6 +161,8 @@ def roll_tensor(tensor, shifts=-1, dims=-1, cp_group=None, packed_seq_params=Non falls back to standard rolling behavior. packed_seq_params (PackedSeqParams): Parameters for packed sequence processing. If provided, respects sequence boundaries. + fill_value: Value to fill at boundary positions where the original sequence has + no data (default 0). For a padding mask, pass ``True``. Returns: tuple: (rolled_tensor, sum_of_rolled_tensor) """ @@ -165,12 +171,14 @@ def roll_tensor(tensor, shifts=-1, dims=-1, cp_group=None, packed_seq_params=Non # Handle packed sequences cases if packed_seq_params is not None: - return _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group) + return _roll_tensor_packed_seq( + tensor, shifts, dims, packed_seq_params, cp_group, fill_value=fill_value + ) # Standard rolling behavior when CP is not enabled (cp_group is None or size=1) if cp_group is None or cp_group.size() == 1: rolled_tensor = torch.roll(tensor, shifts=shifts, dims=dims) - rolled_tensor.select(dims, shifts).fill_(0) + rolled_tensor.select(dims, shifts).fill_(fill_value) return rolled_tensor, rolled_tensor.sum() # CP-enabled rolling: Split tensor into chunks and handle boundary communication @@ -207,8 +215,7 @@ def roll_tensor(tensor, shifts=-1, dims=-1, cp_group=None, packed_seq_params=Non req_recv_second_part = torch.distributed.irecv(tensor=tensor_recv_list[1], src=prev_rank) ops.append(req_recv_second_part) else: - # Inserted elements are set to be 0.0. - tensor_recv_list[1] = 0 + tensor_recv_list[1] = fill_value if local_rank != len(global_ranks) - 1: req_recv_first_part = torch.distributed.irecv(tensor=tensor_recv_list[0], src=next_rank) ops.append(req_recv_first_part) @@ -235,7 +242,9 @@ def roll_tensor(tensor, shifts=-1, dims=-1, cp_group=None, packed_seq_params=Non return rolled_tensor, rolled_tensor.sum() -def _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group=None): +def _roll_tensor_packed_seq( + tensor, shifts, dims, packed_seq_params, cp_group=None, fill_value=0 +): """Roll tensor with packed sequence support. This function handles rolling for packed sequences by respecting sequence boundaries """ @@ -246,24 +255,52 @@ def _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group=No dims == -1 or dims == tensor.dim() - 1 ), "Packed sequence roll only supports the last dimension." assert shifts == -1, "Packed sequence roll only supports a single-token left shift." - cu_seqlens = packed_seq_params.cu_seqlens_q + # Prefer the padded cumulative seqlens because, with CP, the local THD layout is + # produced by `tex.thd_get_partitioned_indices(cu_seqlens_padded, ...)` and requires + # each per-sequence padded length to be divisible by 2*cp_size. Indexing with the + # unpadded cu_seqlens then produces wrong local boundaries when seqlens are not + # already multiples of 2*cp_size (e.g. odd seqlens). + cu_seqlens = ( + packed_seq_params.cu_seqlens_q_padded + if getattr(packed_seq_params, 'cu_seqlens_q_padded', None) is not None + else packed_seq_params.cu_seqlens_q + ) assert cu_seqlens is not None, "Packed sequence parameters must provide cu_seqlens_q." - rolled_tensor = tensor.clone() - cp_size = cp_group.size() if cp_group is not None else 1 if cp_size == 1: + rolled_tensor = tensor.clone() # CP disabled: roll each packed sequence independently within its boundaries for i in range(len(cu_seqlens) - 1): start_idx = cu_seqlens[i] end_idx = cu_seqlens[i + 1] seq_slice = tensor[..., start_idx:end_idx] rolled_seq = torch.roll(seq_slice, shifts=shifts, dims=dims) - # Zero out the last position(s) that would cross sequence boundaries - rolled_seq[..., shifts:] = 0 + rolled_seq[..., shifts:] = fill_value rolled_tensor[..., start_idx:end_idx] = rolled_seq return rolled_tensor, rolled_tensor.sum() + cp_partition_mode = getattr(packed_seq_params, 'cp_partition_mode', 'zigzag') + if cp_partition_mode == 'zigzag': + rolled_tensor = _roll_tensor_packed_seq_zigzag_cp( + tensor, shifts, dims, cu_seqlens, cp_group, fill_value=fill_value + ) + return rolled_tensor, rolled_tensor.sum() + if cp_partition_mode == 'contiguous': + rolled_tensor = _roll_tensor_packed_seq_contiguous_cp( + tensor, dims, cu_seqlens, cp_group, fill_value=fill_value + ) + return rolled_tensor, rolled_tensor.sum() + raise ValueError(f"Unsupported packed sequence CP partition mode: {cp_partition_mode}") + + +def _roll_tensor_packed_seq_zigzag_cp( + tensor, shifts, dims, cu_seqlens, cp_group, fill_value=0 +): + """Roll a zigzag-CP THD shard without crossing packed sequence boundaries.""" + cp_size = cp_group.size() + rolled_tensor = tensor.clone() + # CP enabled: each rank owns two chunks per sequence (front and mirrored tail). local_rank = torch.distributed.get_rank(group=cp_group) global_ranks = torch.distributed.get_process_group_ranks(group=cp_group) @@ -312,7 +349,7 @@ def _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group=No ops.append(torch.distributed.isend(tensor=tensor_send_list[0], dst=prev_rank)) ops.append(torch.distributed.irecv(tensor=tensor_recv_list[1], src=prev_rank)) else: - tensor_recv_list[1].zero_() + tensor_recv_list[1].fill_(fill_value) if local_rank != cp_size - 1: ops.append(torch.distributed.irecv(tensor=tensor_recv_list[0], src=next_rank)) @@ -336,7 +373,56 @@ def _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group=No # update the rolled tensor rolled_tensor[..., local_start_idx:local_end_idx] = seq_result - return rolled_tensor, rolled_tensor.sum() + return rolled_tensor + + +def _roll_tensor_packed_seq_contiguous_cp( + tensor, dims, cu_seqlens, cp_group, fill_value=0 +): + """Roll a contiguous-CP THD shard without crossing packed sequence boundaries.""" + local_seq_len = tensor.size(dims) + rolled_tensor = torch.roll(tensor, shifts=-1, dims=dims) + if local_seq_len == 0: + return rolled_tensor + + cp_size = cp_group.size() + local_rank = torch.distributed.get_rank(group=cp_group) + global_ranks = torch.distributed.get_process_group_ranks(group=cp_group) + + cu = cu_seqlens.to(device=tensor.device, dtype=torch.long) + if cu.numel() > 1: + nonduplicate_boundaries = torch.ones(cu.numel(), device=cu.device, dtype=torch.bool) + nonduplicate_boundaries[1:] = cu[1:] != cu[:-1] + cu = cu[nonduplicate_boundaries] + if cu.numel() <= 1: + rolled_tensor.fill_(fill_value) + return rolled_tensor + + global_start = local_rank * local_seq_len + global_positions = global_start + torch.arange(local_seq_len, device=tensor.device) + seq_idx = torch.bucketize(global_positions, cu[1:], right=True).clamp(max=cu.numel() - 2) + seq_ends = cu[1:][seq_idx] + valid_next = (global_positions < cu[-1]) & (global_positions + 1 < seq_ends) + + rolled_tensor[..., ~valid_next] = fill_value + + recv_next_first = torch.empty_like(tensor.select(dims, 0)) + ops = [] + if local_rank < cp_size - 1: + next_rank = global_ranks[local_rank + 1] + ops.append(torch.distributed.irecv(tensor=recv_next_first, src=next_rank)) + if local_rank > 0: + prev_rank = global_ranks[local_rank - 1] + send_first = tensor.select(dims, 0).contiguous() + ops.append(torch.distributed.isend(tensor=send_first, dst=prev_rank)) + for op in ops: + op.wait() + + if local_rank < cp_size - 1: + last = rolled_tensor.select(dims, -1) + last.copy_(torch.where(valid_next[-1], recv_next_first, last)) + + return rolled_tensor class MTPLossLoggingHelper: @@ -573,6 +659,8 @@ class MultiTokenPredictionLayerSubmodules: layer_norm: LayerNormBuilder eh_proj: Union[ModuleSpec, type] = None + e_proj: Union[ModuleSpec, type] = None + h_proj: Union[ModuleSpec, type] = None mtp_model_layer: Union[ModuleSpec, type] = None @@ -591,7 +679,8 @@ def get_mtp_layer_spec( def get_mtp_layer_spec_for_backend( - mtp_model_layer_spec: ModuleSpec, backend: BackendSpecProvider + mtp_model_layer_spec: ModuleSpec, + backend: BackendSpecProvider, ) -> ModuleSpec: """Get the MTP layer spec. @@ -944,6 +1033,13 @@ def __init__( stacklevel=2, ) hybrid_submodules = mamba_submodules + if self.config.enable_hyper_connections and ( + mtp_layer_pattern is None or hybrid_submodules is None + ): + raise ValueError( + "Multi-token prediction with hyper connections requires the HybridModel " + "MTP contract: both mtp_layer_pattern and hybrid_submodules must be provided." + ) self.sequence_parallel = config.sequence_parallel self.submodules = submodules self.layer_number = layer_number + get_mtp_layer_offset(self.config, vp_stage) @@ -951,6 +1047,7 @@ def __init__( self.cp_group = pg_collection.cp self.tp_group = pg_collection.tp if pg_collection is not None else None self.mtp_layer_pattern = mtp_layer_pattern + self.mhc_enabled = self.config.enable_hyper_connections # Validate attention mask type if using transformer-based inner layers if self.submodules.mtp_model_layer is not None and hasattr( @@ -992,25 +1089,51 @@ def __init__( eps=self.config.layernorm_epsilon, ) - # For the linear projection at the (k - 1)-th MTP layer, the input is the concatenation - # of the i-th token's hidden states and the (i + K)-th token's decoder input, - # so the input's shape is [s, b, 2*h]. - # The output will be send to the following transformer layer, - # so the output's shape should be [s, b, h]. - self.eh_proj = build_module( - self.submodules.eh_proj, - self.config.hidden_size * 2, - self.config.hidden_size, - config=self.config, - init_method=self.config.init_method, - gather_output=False, - bias=False, - skip_bias_add=False, - is_expert=False, - tp_comm_buffer_name="mtp_eh_proj", - tp_group=pg_collection.tp if pg_collection is not None else None, - name=(name + ".eh_proj") if name is not None else None, - ) + if self.mhc_enabled: + projection_kwargs = { + "config": self.config, + "init_method": self.config.init_method, + "gather_output": False, + "bias": False, + "skip_bias_add": False, + "is_expert": False, + "tp_group": pg_collection.tp if pg_collection is not None else None, + } + self.e_proj = build_module( + self.submodules.e_proj, + self.config.hidden_size, + self.config.hidden_size, + tp_comm_buffer_name="mtp_e_proj", + name=(name + ".e_proj") if name is not None else None, + **projection_kwargs, + ) + self.h_proj = build_module( + self.submodules.h_proj, + self.config.hidden_size, + self.config.hidden_size, + tp_comm_buffer_name="mtp_h_proj", + name=(name + ".h_proj") if name is not None else None, + **projection_kwargs, + ) + self.eh_proj = None + else: + # Combine each hidden state with the corresponding future-token embedding. + self.eh_proj = build_module( + self.submodules.eh_proj, + self.config.hidden_size * 2, + self.config.hidden_size, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=False, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name="mtp_eh_proj", + tp_group=pg_collection.tp if pg_collection is not None else None, + name=(name + ".eh_proj") if name is not None else None, + ) + self.e_proj = None + self.h_proj = None # Build inner layers: two possible paths # 1. Hybrid path: use HybridStack for hybrid pattern support @@ -1029,6 +1152,7 @@ def __init__( post_process=True, # MTP layer is self-contained pg_collection=pg_collection, is_mtp_layer=True, + mtp_layer_number=self.layer_number, name=(name + ".mtp_model_layer") if name is not None else None, ) elif self.config.mtp_num_layers is not None: @@ -1051,6 +1175,17 @@ def __init__( hidden_size=self.config.hidden_size, eps=self.config.layernorm_epsilon, ) + if self.mhc_enabled: + hc_mult = self.config.num_residual_streams + hc_dim = self.config.hidden_size * hc_mult + self.hc_head_fn = mark_keep_in_fp32(nn.Parameter(torch.randn(hc_mult, hc_dim))) + self.hc_head_base = mark_keep_in_fp32(nn.Parameter(torch.zeros(hc_mult))) + self.hc_head_scale = mark_keep_in_fp32(nn.Parameter(torch.ones(1))) + nn.init.xavier_uniform_(self.hc_head_fn) + if self.config.sequence_parallel: + setattr(self.hc_head_fn, "sequence_parallel", True) + setattr(self.hc_head_base, "sequence_parallel", True) + setattr(self.hc_head_scale, "sequence_parallel", True) self.offload_context = nullcontext() def _get_embeddings( @@ -1099,6 +1234,7 @@ def _get_embeddings( dims=-1, cp_group=self.cp_group, packed_seq_params=packed_seq_params, + fill_value=True, ) # embedding decoder_input = embedding(input_ids=input_ids, position_ids=position_ids) @@ -1123,25 +1259,48 @@ def _concat_embeddings(self, hidden_states: torch.Tensor, decoder_input: torch.T """ decoder_input = apply_module(self.enorm)(decoder_input) decoder_input = make_viewless_tensor(inp=decoder_input, requires_grad=True, keep_graph=True) - hidden_states = apply_module(self.hnorm)(hidden_states) - hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) - # At the (k - 1)-th MTP module, concatenates the i-th token's hidden_states - # and the (i + K)-th token's embedding, and combine them with linear projection. - hidden_states = torch.cat((decoder_input, hidden_states), -1) - hidden_states, _ = self.eh_proj(hidden_states) - # For tensor parallel we need to gather the tensor across the model-parallel - # ranks after the linear projection. - if InferenceMode.is_active(): - hidden_states = inference_all_gather_from_tensor_model_parallel_region( - hidden_states, self.tp_group, self.config + + if self.mhc_enabled: + n = self.config.num_residual_streams + h = self.config.hidden_size + seq_len, batch_size, _ = hidden_states.shape + hidden_streams = hidden_states.view(seq_len, batch_size, n, h) + hidden_streams = apply_module(self.hnorm)(hidden_streams) + hidden_streams = make_viewless_tensor( + inp=hidden_streams, requires_grad=True, keep_graph=True + ) + embedded, _ = self.e_proj(decoder_input) + embedded = gather_from_tensor_model_parallel_region(embedded, group=self.tp_group) + projected_hidden, _ = self.h_proj(hidden_streams) + projected_hidden = gather_from_tensor_model_parallel_region( + projected_hidden, group=self.tp_group ) + seq_len, batch_size, n, h = projected_hidden.shape + embedded = embedded.unsqueeze(2).expand(seq_len, batch_size, n, h) + hidden_states = (embedded + projected_hidden).reshape(seq_len, batch_size, n * h) + if self.sequence_parallel: + hidden_states = scatter_to_sequence_parallel_region( + hidden_states, group=self.tp_group + ) else: - hidden_states = gather_from_tensor_model_parallel_region( - hidden_states, group=self.tp_group + hidden_states = apply_module(self.hnorm)(hidden_states) + hidden_states = make_viewless_tensor( + inp=hidden_states, requires_grad=True, keep_graph=True ) - # For sequence parallel, scatter after linear_fc and before transformer layer. - if self.sequence_parallel: - hidden_states = scatter_to_sequence_parallel_region(hidden_states, group=self.tp_group) + hidden_states = torch.cat((decoder_input, hidden_states), -1) + hidden_states, _ = self.eh_proj(hidden_states) + if InferenceMode.is_active(): + hidden_states = inference_all_gather_from_tensor_model_parallel_region( + hidden_states, self.tp_group, self.config + ) + else: + hidden_states = gather_from_tensor_model_parallel_region( + hidden_states, group=self.tp_group + ) + if self.sequence_parallel: + hidden_states = scatter_to_sequence_parallel_region( + hidden_states, group=self.tp_group + ) return hidden_states def _proj_and_transformer_layer( @@ -1212,7 +1371,8 @@ def _proj_and_transformer_layer( padding_mask=padding_mask, ) - hidden_states = self._postprocess(hidden_states) + if not self.mhc_enabled: + hidden_states = self._postprocess(hidden_states) return hidden_states @@ -1221,6 +1381,16 @@ def _postprocess(self, hidden_states: torch.Tensor): Postprocesses the output of the transformer layers. """ + if self.mhc_enabled: + hidden_states = learned_output_contract( + hidden_states, + self.hc_head_fn, + self.hc_head_base, + self.hc_head_scale, + self.config.num_residual_streams, + self.config.layernorm_epsilon, + ) + # Layer norm before shared head layer. hidden_states = apply_module(self.final_layernorm)(hidden_states) # TENorm produces a "viewed" tensor. This will result in schedule.py's @@ -1793,6 +1963,7 @@ def forward( sequence_len_offset: Optional[Tensor] = None, extra_block_kwargs: Optional[dict] = None, embedding=None, + mhc_multistream: Optional[Tensor] = None, ) -> Tensor: """ Perform the forward pass through all of the MTP modules. @@ -1800,6 +1971,8 @@ def forward( Args: hidden_states (Tensor): Hidden states for input token with the shape [s, b, h] where s is the sequence length, b is the batch size, and h is the hidden size. + mhc_multistream (Tensor, optional): Pre-contraction decoder output [s, b, n*h] + used as the input to MTP depths when hyper connections are enabled. attention_mask (Tensor): Boolean tensor of shape [1, 1, s, s] for masking self-attention. @@ -1809,7 +1982,11 @@ def forward( # get hidden states from previous mtp stages offset = get_mtp_layer_offset(self.config, self.vp_stage) hidden_states_list = list(torch.chunk(hidden_states, 1 + offset, dim=0)) - hidden_states = hidden_states_list[offset] + if mhc_multistream is not None: + mhc_chunks = list(torch.chunk(mhc_multistream, 1 + offset, dim=0)) + hidden_states = mhc_chunks[offset] + else: + hidden_states = hidden_states_list[offset] if self.config.mtp_detach_heads: hidden_states = hidden_states.detach() @@ -1832,9 +2009,11 @@ def forward( **(extra_block_kwargs or {}), ) - # append the output hidden states of the current mtp layer - # to the hidden_states_list - hidden_states_list.append(hidden_states) + if mhc_multistream is not None: + mhc_chunks.append(hidden_states) + hidden_states_list.append(self.layers[layer_idx]._postprocess(hidden_states)) + else: + hidden_states_list.append(hidden_states) # concat the hidden states of all mtp layers hidden_states = torch.cat(hidden_states_list, dim=0) diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py index 0415035ffbe..732adbbebd1 100755 --- a/megatron/core/transformer/transformer_block.py +++ b/megatron/core/transformer/transformer_block.py @@ -1,8 +1,9 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + import logging from contextlib import nullcontext from dataclasses import dataclass -from typing import List, Optional, Set, Union, cast +from typing import List, Optional, Set, Tuple, Union, cast import torch from torch import Tensor @@ -21,8 +22,10 @@ from megatron.core.pipeline_parallel.utils import is_vp_first_stage, is_vp_last_stage from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.recompute import checkpointed_forward +from megatron.core.tensor_parallel.random import CheckpointWithoutOutputManager from megatron.core.transformer.cuda_graphs import annotate_first_last_layer from megatron.core.transformer.enums import InferenceCudaGraphScope, LayerType +from megatron.core.transformer.hyper_connection import HyperConnectionModule from megatron.core.transformer.module import GraphableMegatronModule, MegatronModule from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.torch_norm import LayerNormBuilder @@ -318,6 +321,7 @@ def __init__( self.offload_context, self.group_prefetch_offload_commit_async = nullcontext(), None self.config._cpu_offloading_context = None + self.num_residual_streams = config.num_residual_streams self._build_layers() self.num_layers_per_pipeline_rank = len(self.layers) @@ -483,6 +487,46 @@ def __call__(self, *args, **kwargs): return super().__call__(*args, **kwargs)[0] return super().__call__(*args, **kwargs) + def _build_mhc_recompute_layer_plan( + self, use_mhc_recompute: bool + ) -> Tuple[List[Optional[CheckpointWithoutOutputManager]], List[bool]]: + """Pre-build per-layer MHC recompute managers and block-end markers.""" + num_layers = len(self.layers) + layer_managers: List[Optional[CheckpointWithoutOutputManager]] = [None] * num_layers + is_recompute_block_end: List[bool] = [False] * num_layers + + if not use_mhc_recompute or num_layers == 0: + return layer_managers, is_recompute_block_end + + mhc_recompute_layer_num = self.config.mhc_recompute_layer_num + mhc_manager = CheckpointWithoutOutputManager() + + for l_no in range(num_layers): + is_last_in_transformer_block = l_no == num_layers - 1 + is_last_in_recompute_block = is_last_in_transformer_block + if mhc_recompute_layer_num is not None: + is_last_in_recompute_block = is_last_in_transformer_block or ( + (l_no + 1) % mhc_recompute_layer_num == 0 + ) + + layer_managers[l_no] = mhc_manager + is_recompute_block_end[l_no] = is_last_in_recompute_block + + if is_last_in_recompute_block and not is_last_in_transformer_block: + mhc_manager = CheckpointWithoutOutputManager() + + return layer_managers, is_recompute_block_end + + @staticmethod + def _finalize_mhc_recompute_layer( + mhc_manager: Optional[CheckpointWithoutOutputManager], + hidden_states: Tensor, + is_last_in_recompute_block: bool, + ) -> None: + """Finalize MHC recompute state for the current layer when block ends.""" + if mhc_manager is not None and is_last_in_recompute_block: + mhc_manager.discard_all_outputs_and_register_unified_recompute(hidden_states) + def forward( self, hidden_states: Union[Tensor, WrappedTensor], @@ -592,6 +636,13 @@ def forward( # is called here to be future-proof and corner-case-proof. hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) + # Expand hidden states for hyper connections at the start of the block + # Only expand at the first PP stage; subsequent stages receive n-stream from previous stage + if self.config.enable_hyper_connections and self.pre_process: + hidden_states = HyperConnectionModule.input_expand( + hidden_states, self.num_residual_streams + ) # [s, b, C] -> [s, b, n*C] + if self.config.sequence_parallel: rng_context = tensor_parallel.get_cuda_rng_tracker().fork() else: @@ -619,6 +670,18 @@ def forward( use_inner_quantization_context = False outer_quantization_context = nullcontext() + # Determine if MHC recompute should be used + # Only enable when: training mode AND hyper connections AND 'mhc' in recompute_modules + use_mhc_recompute = ( + self.training + and self.config.enable_hyper_connections + and self.config.recompute_granularity == 'selective' + and "mhc" in self.config.recompute_modules + ) + mhc_layer_managers, mhc_is_last_in_recompute_block = self._build_mhc_recompute_layer_plan( + use_mhc_recompute + ) + with rng_context, outer_quantization_context: # Forward pass. if self.config.recompute_granularity == 'full' and self.training: @@ -660,6 +723,19 @@ def forward( else: inner_quantization_context = nullcontext() + mhc_manager = mhc_layer_managers[l_no] + if mhc_manager is not None: + mhc_manager.is_last_layer_in_recompute_block = ( + mhc_is_last_in_recompute_block[l_no] + ) + + # Only thread mhc_recompute_manager when the layer is mHC and a + # manager actually exists. Plain TransformerLayer (and its + # MoETransformerLayer subclass) doesn't accept this kwarg, and + # its CUDA-graph machinery rejects unrecognized non-tensor kwargs. + extra_layer_kwargs = ( + {"mhc_recompute_manager": mhc_manager} if mhc_manager is not None else {} + ) with self.offload_context, inner_quantization_context: hidden_states, context = layer( hidden_states=hidden_states, @@ -675,7 +751,13 @@ def forward( packed_seq_params=packed_seq_params, sequence_len_offset=sequence_len_offset, padding_mask=padding_mask, + **extra_layer_kwargs, ) + self._finalize_mhc_recompute_layer( + mhc_manager=mhc_manager, + hidden_states=hidden_states, + is_last_in_recompute_block=mhc_is_last_in_recompute_block[l_no], + ) if ( torch.is_grad_enabled() @@ -688,6 +770,12 @@ def forward( if (l_no + layer_offset) in extract_layer_indices: intermediate_hidden_states.append(hidden_states) + # Only contract if the final layer norm is in this stage + if self.config.enable_hyper_connections and self.has_final_layernorm_in_this_stage(): + hidden_states = HyperConnectionModule.output_contract( + hidden_states, self.num_residual_streams + ) # [s, b, n*C] -> [s, b, C] + # Final layer norm. if self.final_layernorm is not None: hidden_states = apply_module(self.final_layernorm)(cast(Tensor, hidden_states)) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 0c9ce022db7..3f24572f101 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 @@ -24,6 +25,7 @@ CudaGraphModule, CudaGraphScope, InferenceCudaGraphScope, + LayerType, ) from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout @@ -217,7 +219,7 @@ class TransformerConfig(ModelParallelConfig): activation_func_clamp_value: Optional[float] = None """Clamp the output of the linear_fc1 in the activation function. Only used when activation_func - is quick_gelu.""" + is quick_gelu or SwiGLU (MoE only).""" num_moe_experts: Optional[int] = None """Number of experts to use for MoE layer. When set, it replaces MLP with MoE layer. Set to None @@ -281,12 +283,17 @@ class TransformerConfig(ModelParallelConfig): #################### # attention variant #################### - experimental_attention_variant: Optional[Literal['gated_delta_net', 'dsa']] = None - """Type of attention variant to use. Currently support gated_delta_net and dsa.""" + experimental_attention_variant: Optional[Literal['gated_delta_net', 'dsa', 'dsv4_hybrid']] = ( + None + ) + """Type of attention variant to use. Currently support gated_delta_net, dsa, and dsv4_hybrid.""" experimental_attention_variant_loss_scale_func: Optional[Callable[[torch.Tensor], None]] = None """Optional hook for experimental attention variants to receive the main loss scale.""" + cp_partition_mode: Literal["zigzag", "contiguous"] = "zigzag" + """How THD sequence rows are partitioned across context-parallel ranks.""" + #################### # DSA #################### @@ -333,6 +340,27 @@ class TransformerConfig(ModelParallelConfig): dsa_indexer_k_norm_fp32: bool = False """Whether DSA indexer key LayerNorm should run on fp32 inputs.""" + #################### + # Compressed sparse attention + #################### + csa_window_size: int = 128 + """Sliding window size for compressed sparse attention.""" + + csa_compress_ratios: Optional[List[int]] = None + """Per-layer compress ratios, e.g. [0, 0, 4, 128, 4, 128, ...].""" + + csa_compress_rotary_base: float = 40000.0 + """RoPE base for compressed KV positions in compressed sparse attention.""" + + csa_dense_mode: bool = False + """Whether to use dense mode for compressed sparse attention. If True, the CSA indexer will be + disabled.""" + + apply_dsa_kernel_fusion: bool = False + """If True, use fused DSA sparse-attention kernels (FlashMLA forward + cuDNN DSA backward, + indexer scoring, and top-K selection). Requires ``flash_mla`` and ``nvidia-cudnn-frontend`` + with CuTe-DSL support. When False, falls back to unfused PyTorch implementations.""" + #################### # linear attention #################### @@ -536,7 +564,7 @@ class TransformerConfig(ModelParallelConfig): recompute_modules: Optional[List[str]] = None """The submodules to recompute. choices: "core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", - "shared_experts", "gdn_norm_out". + "shared_experts", "gdn_norm_out", "mhc". default: ["core_attn"]. "core_attn": recompute the core attention part of the transformer layer. "moe_act": recompute the MoE MLP activation function. @@ -546,7 +574,11 @@ class TransformerConfig(ModelParallelConfig): "moe": recompute the MoE layer. "shared_experts": recompute the shared experts in the MoE layer. "gdn_norm_out": recompute the GatedDeltaNet output norm and HP-to-CP all-to-all. - "moe_act", "layernorm", "mla_up_proj", and "gdn_norm_out" use output-discarding checkpointing, + "mhc": recompute HyperConnection intermediate activations via + CheckpointWithoutOutput + CheckpointWithoutOutputManager. Requires + enable_hyper_connections=True. Cannot be used with "mlp". + "moe_act", "layernorm", "mla_up_proj", "gdn_norm_out", and "mhc" use + output-discarding checkpointing, "core_attn", "mlp", "moe", and "shared_experts" use normal checkpointing. """ @@ -810,6 +842,15 @@ class TransformerConfig(ModelParallelConfig): If negative, generates bias once per layer and reuses it (abs value is std). This is an experimental feature for benchmarking purposes.""" + moe_n_hash_layers: int = 0 + """Number of leading transformer layers that use hash-based MoE routing. + Layers with ``layer_number <= moe_n_hash_layers`` select experts from a + token-to-expert lookup table instead of learned top-k routing.""" + + actual_vocab_size: Optional[int] = None + """Vocabulary size of the token-to-expert lookup table. + Required when ``moe_n_hash_layers > 0``.""" + use_grouped_gemm_for_dense_mlp: bool = False """Use GroupedLinear(num_groups=1) for dense MLP to trigger the ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 fusion on SM100+ with MXFP8 recipe. @@ -865,6 +906,13 @@ class TransformerConfig(ModelParallelConfig): moe_permute_fusion_into_hybridep: bool = False """Fuse token rearrangement ops during token dispatching for HybridEP.""" + moe_hybridep_pad_uneven_dispatch_inputs: bool = False + """Pad uneven HybridEP dispatch inputs to the group maximum before dispatch. + Enable when local HybridEP input token counts can differ across ranks, for example + with dynamically packed THD inputs. Leave disabled when dispatcher inputs are + already padded to equal token counts. + """ + moe_per_layer_logging: bool = False """Enable per-layer logging for MoE, currently supports auxiliary loss and z loss.""" @@ -1043,6 +1091,17 @@ class TransformerConfig(ModelParallelConfig): transformed to an empty list in __post_init__. The deprecated values "full_iteration" and "full_iteration_inference" are also accepted and migrated to the new API in __post_init__.""" + create_attention_mask_in_dataloader: bool = True + """Whether data loaders create and pass an explicit attention mask.""" + + thd_max_packed_sequences: int = field( + default=32, metadata={"argparse_meta": {"arg_names": ["--thd-max-packed-sequences"]}} + ) + """Static THD sequence capacity, including an optional dummy padding sequence.""" + + cuda_graph_dynamic_microbatches: bool = False + """Capture enough TE graph slots for varying packed-microbatch counts.""" + inference_cuda_graph_scope: Optional[InferenceCudaGraphScope] = field( default=None, metadata={ @@ -1074,6 +1133,42 @@ class TransformerConfig(ModelParallelConfig): CudaGraphScope instances deserialized from pre-refactor checkpoints are converted to their string names before normalization so existing CUDA_GRAPH_MODULES_DEPRECATIONS handles them.""" + #################### + # Hyper-Connection Configuration + #################### + enable_hyper_connections: bool = False + """Enable mHC residual connections.""" + + num_residual_streams: int = 4 + """Number of residual streams (n in paper).""" + + mhc_sinkhorn_iterations: int = 20 + """Number of Sinkhorn-Knopp iterations for doubly stochastic projection.""" + + mhc_init_gating_factor: float = 0.01 + """Initial value of Gating Factor (alpha in paper).""" + + use_fused_mhc: bool = False + """Use fused kernels for mHC operations when supported. + + Backend selection is operation-specific, with native torch fallbacks that + preserve the same public behavior when Triton or cuTile is unavailable. + """ + + mhc_recompute_layer_num: Optional[int] = None + """Number of layers per MHC recompute block. + + When set, every `mhc_recompute_layer_num` layers form a recompute block. The last layer + in each recompute block (i.e., layer_number % mhc_recompute_layer_num == 0 or the final + layer in the transformer block) will: + - NOT checkpoint its final MLP BDA + - Register the unified recompute hook on its MLP BDA output + - A new CheckpointWithoutOutputManager is created for subsequent layers + + If None, all layers in the transformer block share a single recompute block. + + Must be a positive integer when set.""" + #################### # miscellaneous #################### @@ -1325,11 +1420,44 @@ def __post_init__(self): f"tensor_model_parallel_size ({self.tensor_model_parallel_size})." ) + 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.hybrid_context_parallel + ): + raise ValueError( + "DSv4 Hybrid does not support dynamic per-microbatch context parallelism; " + "use a static context_parallel_size." + ) + + if self.context_parallel_size > 1: + if self.experimental_attention_variant == "dsv4_hybrid": + if self.sequence_packing_scheduler is None: + raise ValueError( + "DSv4 Hybrid with CP requires a sequence_packing_scheduler for THD inputs." + ) + if self.cp_partition_mode != "contiguous": + raise ValueError( + "DSv4 Hybrid with CP requires cp_partition_mode='contiguous'." + ) + elif self.cp_partition_mode != "zigzag": + raise ValueError( + "cp_partition_mode='contiguous' currently is only supported with dsv4_hybrid." + ) + if self.experimental_attention_variant == "gated_delta_net": assert ( self.linear_attention_freq is not None ), f"linear_attention_freq must be set for linear gated_delta_net." + if self.pad_packed_seq_alignment is not None: + assert self.pad_packed_seq_by_appending_dummy_seq, ( + "gated_delta_net with pad_packed_seq_alignment requires " + "pad_packed_seq_by_appending_dummy_seq." + ) + # Check required parameters assert ( self.linear_conv_kernel_dim is not None @@ -1377,6 +1505,101 @@ def __post_init__(self): "dsa_indexer_skip_topk_offset must be non-negative, got " f"{self.dsa_indexer_skip_topk_offset}." ) + elif self.experimental_attention_variant == "dsv4_hybrid": + assert self.multi_latent_attention, "DSv4 Hybrid requires multi_latent_attention." + assert self.csa_compress_ratios is not None, "csa_compress_ratios must be set" + mtp_layers = self.mtp_num_layers or 0 + expected_len = self.num_layers + mtp_layers + # A HybridModel MTP depth may contain more than one hybrid layer, while + # mtp_num_layers counts depths. The per-layer ratio list therefore has a minimum, + # rather than an exact, length here; HybridModel argument normalization checks its + # exact pattern-derived length before constructing this config. + assert len(self.csa_compress_ratios) >= expected_len, ( + f"csa_compress_ratios length ({len(self.csa_compress_ratios)}) must be at least " + f"num_layers + mtp_num_layers ({self.num_layers} + {mtp_layers} = {expected_len})" + ) + assert all( + ratio in [0, 4, 128] for ratio in self.csa_compress_ratios + ), "csa_compress_ratios must be 0, 4, or 128" + assert ( + self.tensor_model_parallel_size == 1 + ), "DSv4 Hybrid Attention only supports TP size 1." + assert not self.qk_clip, "QK clipping is not supported with DSv4 Hybrid Attention." + self.hetereogenous_dist_checkpoint = True + + if self.apply_dsa_kernel_fusion: + assert ( + 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] >= 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: + from flash_mla import flash_mla_sparse_fwd # noqa: F401 + except ImportError: + flash_mla_available = False + + cudnn_dsa_available = True + try: + from cudnn import DSA # noqa: F401 + except ImportError: + cudnn_dsa_available = False + + if not flash_mla_available or not cudnn_dsa_available: + missing = [] + if not flash_mla_available: + missing.append( + "flash_mla (install from " + "https://github.com/deepseek-ai/FlashMLA/tree/nv_dev)" + ) + if not cudnn_dsa_available: + missing.append("cudnn-frontend DSA (nvidia-cudnn-frontend[cutedsl])") + raise ValueError( + f"apply_dsa_kernel_fusion requires fused DSA kernels, but the " + f"following packages are not available: {', '.join(missing)}. " + f"Install them or pass --no-dsa-kernel-fusion to use the unfused " + f"PyTorch fallback." + ) + + if self.context_parallel_size > 1 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.fp8: # cannot support first last layer bf16 with delayed scaling @@ -1755,6 +1978,7 @@ def __post_init__(self): "moe", "shared_experts", "gdn_norm_out", + "mhc", } invalid_modules = set(self.recompute_modules) - allowed_modules assert not invalid_modules, ( @@ -1826,6 +2050,53 @@ def __post_init__(self): if "moe" not in self.recompute_modules: self.recompute_modules.append("moe") + # Validation for "mhc" in recompute_modules + if self.recompute_granularity == "selective" and "mhc" in self.recompute_modules: + if not self.enable_hyper_connections: + raise ValueError( + "'mhc' in recompute_modules requires enable_hyper_connections=True." + ) + if "mlp" in self.recompute_modules: + raise ValueError( + "'mhc' and 'mlp' in recompute_modules cannot be used together. " + "They use different checkpoint mechanisms that may conflict." + ) + if self.mhc_recompute_layer_num is not None and ( + isinstance(self.mhc_recompute_layer_num, bool) + or not isinstance(self.mhc_recompute_layer_num, int) + or self.mhc_recompute_layer_num < 1 + ): + raise ValueError( + "mhc_recompute_layer_num must be a positive integer when " + "'mhc' is in recompute_modules." + ) + if self.fine_grained_activation_offloading: + raise NotImplementedError( + "'mhc' in recompute_modules + fine_grained_activation_offloading is " + "not yet supported. The mHC recompute hook currently fires before " + "the offloading backward chunk is initialized, causing tensor_pop " + "on a None chunk. Disable one of them." + ) + + if ( + self.enable_hyper_connections + and not self.fine_grained_activation_offloading + and self.cuda_graph_impl != "transformer_engine" + and not self.external_cuda_graph + and not ( + self.recompute_granularity == "selective" + and "mhc" in self.recompute_modules + ) + ): + warnings.warn( + "HyperConnections are enabled but 'mhc' is not in " + "recompute_modules with selective recompute. Consider adding 'mhc' to " + "recompute_modules with selective recompute to reduce activation memory." + ) + + if self.use_fused_mhc and not self.enable_hyper_connections: + raise ValueError("use_fused_mhc requires enable_hyper_connections=True.") + if self.fine_grained_activation_offloading: assert ( not self.cpu_offloading @@ -2143,6 +2414,33 @@ def __post_init__(self): if self.activation_func != F.silu or not self.gated_linear_unit: raise ValueError("Storing activation input in FP8 is supported only for SwiGLU.") + if ( + self.activation_func_clamp_value is not None + and self.activation_func == F.silu + and self.gated_linear_unit + ): + if ( + not math.isfinite(self.activation_func_clamp_value) + or self.activation_func_clamp_value <= 0 + ): + raise ValueError( + "activation_func_clamp_value for SwiGLU must be finite and greater than zero." + ) + if self.num_moe_experts is None: + raise ValueError( + "activation_func_clamp_value for SwiGLU is only supported with MoE." + ) + if self.glu_linear_offset != 0.0: + raise ValueError( + "glu_linear_offset must be zero when activation_func_clamp_value " + "is set for SwiGLU." + ) + if self.use_te_activation_func: + raise ValueError( + "use_te_activation_func must be False " + "when activation_func_clamp_value is not None for SwiGLU" + ) + if self.apply_rope_fusion: if self.multi_latent_attention: warnings.warn( @@ -2272,6 +2570,37 @@ def __post_init__(self): "'sqrtsoftplus', or unset --moe-router-enable-expert-bias." ) + if self.moe_n_hash_layers > 0: + assert ( + self.actual_vocab_size is not None and self.actual_vocab_size > 0 + ), "actual_vocab_size must be positive when moe_n_hash_layers > 0." + assert ( + self.num_moe_experts is not None + ), "num_moe_experts must be set when moe_n_hash_layers > 0." + if self.pipeline_model_parallel_size > 1 and not self.is_hybrid_model: + assert self.pipeline_model_parallel_layout is not None, ( + "pipeline_model_parallel_layout must be set when using hash MoE " + "layers with pipeline parallelism (PP > 1)." + ) + embedding_stage = self.pipeline_model_parallel_layout.layout[0][0] + n_decoders_with_embedding = embedding_stage.count(LayerType.decoder) + assert self.moe_n_hash_layers <= n_decoders_with_embedding, ( + "All hash MoE layers must currently share the virtual pipeline stage " + "that owns the embedding. The embedding stage has " + f"{n_decoders_with_embedding} decoder layers, but " + f"moe_n_hash_layers={self.moe_n_hash_layers}." + ) + assert ( + not self.overlap_moe_expert_parallel_comm + ), "overlap_moe_expert_parallel_comm does not support hash MoE layers yet." + log_single_rank( + logger, + logging.WARNING, + "Hash MoE initialized with a placeholder round-robin token-to-expert table. " + "Load a trained table from a checkpoint or provide a workload-aware " + "initialization before training.", + ) + if self.num_moe_experts and self.fp8: # TE version below 1.7.0 will raise Error when handle zeros tokens for expert if not is_te_min_version("1.7.0.dev0"): @@ -2429,6 +2758,22 @@ def _scope_to_str(s): "local", "full_iteration", ], f"Invalid cuda graph implementation: {self.cuda_graph_impl}" + if ( + self.cuda_graph_impl == "transformer_engine" + and self.recompute_granularity == "selective" + and "mhc" in self.recompute_modules + ): + raise NotImplementedError( + "'mhc' in recompute_modules is not supported with " + "cuda_graph_impl='transformer_engine'. TE CUDA graph replay bypasses " + "the per-forward mHC recompute manager. Remove 'mhc' from " + "recompute_modules when using TE CUDA graphs." + ) + if self.cuda_graph_impl != "transformer_engine": + assert not self.cuda_graph_dynamic_microbatches, ( + "cuda_graph_dynamic_microbatches is only supported with " + "cuda_graph_impl=transformer_engine." + ) self.inference_cuda_graph_scope = normalize_inference_cuda_graph_scope( self.inference_cuda_graph_scope, self.cuda_graph_impl @@ -2829,6 +3174,65 @@ def _scope_to_str(s): self.attention_backend == AttnBackend.flash ), "Batch invariant mode only supports FlashAttention" + if ( + self.cuda_graph_impl == "transformer_engine" + and self.sequence_packing_scheduler is not None + ): + assert self.pad_packed_seq_alignment is not None, ( + "THD CUDA Graph requires pad_packed_seq_alignment to be set." + ) + assert ( + self.pad_packed_seq_alignment == "max" + or self.pad_packed_seq_alignment == self.max_seqlen_per_dp_cp_rank + ), ( + "THD CUDA Graph requires pad_packed_seq_alignment='max' or an alignment " + "equal to max_seqlen_per_dp_cp_rank " + f"({self.max_seqlen_per_dp_cp_rank}), got " + f"{self.pad_packed_seq_alignment}." + ) + if ( + self.cuda_graph_impl == "transformer_engine" + and self.num_moe_experts is not None + and self.sequence_parallel + and self.tensor_model_parallel_size > 1 + ): + raise ValueError( + "THD Transformer Engine CUDA graphs with MoE do not yet support " + "tensor_parallel_size > 1 with sequence_parallel: runtime " + "padding_mask/input_ids remain at the full CP-local length while " + "router activations are TP-scattered." + ) + + if self.sequence_packing_scheduler is not None: + if not HAVE_PACKAGING: + raise ImportError( + "packaging is not installed. Please install it with `pip install packaging`." + ) + if not ( + is_te_min_version("2.9.0") or get_te_version() == PkgVersion("2.9.0.dev0+5b3092a") + ): + raise ValueError( + "THD sequence packing requires Transformer Engine >= 2.9.0 " + f"but got {get_te_version()} (TE < 2.9.0 may have convergence issues)." + ) + + self.variable_seq_lengths = True + assert self.num_moe_experts is None or self.moe_token_dispatcher_type in ( + "alltoall", + "flex", + ), ( + "sequence_packing only supports moe_token_dispatcher_type in " + "('alltoall', 'flex'), " + f"got '{self.moe_token_dispatcher_type}'" + ) + + supported_schedulers = ['dp_balanced'] + if self.sequence_packing_scheduler not in supported_schedulers: + raise ValueError( + f"Unsupported scheduler: {self.sequence_packing_scheduler}. " + f"Available schedulers: {supported_schedulers}" + ) + @dataclass class MLATransformerConfig(TransformerConfig): @@ -2845,10 +3249,12 @@ class MLATransformerConfig(TransformerConfig): """Rank of Query tensor's low rank representation.""" kv_lora_rank: int = 512 - """Rank of Key and Value tensors' low rank representation.""" + """Rank of Key and Value tensors' low rank representation. + This is not used for DSv4 Hybrid Attention and will be overridden automatically.""" qk_head_dim: int = 128 - """Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim""" + """Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim + This is not used for DSv4 Hybrid Attention and will be overridden automatically.""" qk_pos_emb_head_dim: int = 64 """Dimension of the position embedding in the QK projection.""" @@ -2886,6 +3292,12 @@ class MLATransformerConfig(TransformerConfig): mscale_all_dim: float = 0.0 """Mscale all dimensions for YaRN RoPE in Multi-Latent Attention, used by yarn.""" + o_groups: int = 8 + """Number of groups for grouped low-rank output projection (wo_a).""" + + o_lora_rank: int = 1024 + """Low-rank dimension per group for grouped output (wo_a). Used when o_groups > 0.""" + cache_mla_latents: bool = False """Cache the low dimensional tensors for MLA rather than full KV cache. This is only for the dynamic inference backend and requires that @@ -2898,12 +3310,41 @@ class MLATransformerConfig(TransformerConfig): def __post_init__(self): super().__post_init__() - if self.multi_latent_attention and self.apply_rope_fusion and self.rope_type != "yarn": + if ( + self.multi_latent_attention + and self.apply_rope_fusion + and self.rope_type != "yarn" + and self.experimental_attention_variant != "dsv4_hybrid" + ): raise ValueError("apply_rope_fusion for MLA only works with YARN RoPE.") if self.attention_output_gate: raise NotImplementedError("Output gate is not supported for MLA yet.") + # DSv4 hybrid: derive qk_head_dim and kv_lora_rank from v_head_dim and qk_pos_emb_head_dim. + if self.experimental_attention_variant == "dsv4_hybrid": + assert ( + not self.mla_down_proj_fusion + ), "MLA down projection fusion must be disabled for DSv4 hybrid mode." + assert self.q_lora_rank is not None, "DSv4 hybrid mode requires q_lora_rank." + assert self.o_groups > 0, "DSv4 hybrid mode requires o_groups to be positive." + assert self.o_lora_rank > 0, "DSv4 hybrid mode requires o_lora_rank to be positive." + assert ( + self.num_attention_heads * self.v_head_dim + ) % self.o_groups == 0, ( + "num_attention_heads * v_head_dim must be divisible by o_groups." + ) + log_single_rank( + logger, + logging.WARNING, + "DSv4 hybrid mode is enabled, deriving qk_head_dim and kv_lora_rank from " + "v_head_dim and qk_pos_emb_head_dim", + ) + derived = self.v_head_dim - self.qk_pos_emb_head_dim + assert derived > 0, "v_head_dim must be greater than qk_pos_emb_head_dim." + self.qk_head_dim = derived + self.kv_lora_rank = derived + if self.cache_mla_latents: assert ( self.apply_rope_fusion is False diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index f6ea382077e..c165a375afd 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -8,6 +8,9 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Dict, Optional, Protocol, Union +if TYPE_CHECKING: + from megatron.core.tensor_parallel.random import CheckpointWithoutOutputManager + import torch import torch.distributed from torch import Tensor @@ -16,10 +19,21 @@ from megatron.core.dist_checkpointing.mapping import ShardedStateDict from megatron.core.dist_checkpointing.utils import apply_prefix_mapping from megatron.core.inference.utils import InferenceMode -from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.packed_seq_params import ( + CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX, + PackedSeqParams, + build_packed_seq_params_from_cuda_graph_kwargs, + has_packed_seq_params_cuda_graph_kwargs, + split_packed_seq_params_for_cuda_graph, +) from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.cuda_graphs import is_graph_capturing, is_graph_warmup, make_weakref -from megatron.core.transformer.enums import CudaGraphModule, InferenceCudaGraphScope, LayerType +from megatron.core.transformer.enums import ( + AttnMaskType, + CudaGraphModule, + InferenceCudaGraphScope, + LayerType, +) from megatron.core.transformer.identity_op import IdentityFuncOp, IdentityOp from megatron.core.transformer.mlp import MLP from megatron.core.transformer.module import GraphableMegatronModule @@ -268,14 +282,17 @@ class TransformerLayerSubmodules: """ input_layernorm: LayerNormBuilder = IdentityOp + self_attention_hyper_connection: Union[ModuleSpec, type] = IdentityOp self_attention: Union[ModuleSpec, type] = IdentityOp self_attn_bda: Union[ModuleSpec, type] = IdentityFuncOp pre_cross_attn_layernorm: LayerNormBuilder = IdentityOp + cross_attention_hyper_connection: Union[ModuleSpec, type] = IdentityOp cross_attention: Union[ModuleSpec, type] = IdentityOp cross_attn_bda: Union[ModuleSpec, type] = IdentityFuncOp pre_mlp_layernorm: LayerNormBuilder = IdentityOp + mlp_hyper_connection: Union[ModuleSpec, type] = IdentityOp mlp: MlpBuilder | type[IdentityOp] = IdentityOp mlp_bda: Union[ModuleSpec, type] = IdentityFuncOp @@ -368,6 +385,8 @@ def __init__( attention_optional_kwargs["pg_collection"] = pg_collection if pp_layer_offset is not None: attention_optional_kwargs["pp_layer_offset"] = pp_layer_offset + if is_mtp_layer: + attention_optional_kwargs["is_mtp_layer"] = True # [Module 2: SelfAttention] self.self_attention = build_module( @@ -566,6 +585,91 @@ def _get_layer_offset(config: TransformerConfig): ) return get_transformer_layer_offset(config) + @staticmethod + def _group_offload_output_with_bias( + output_with_bias, offload_manager, forced_released_tensors: Optional[list[Tensor]] = None + ): + """Commit a fine-grained offload group for a raw branch output tuple.""" + if isinstance(output_with_bias, tuple): + output = offload_manager.group_offload( + output_with_bias[0], forced_released_tensors=forced_released_tensors + ) + return (output, *output_with_bias[1:]) + return offload_manager.group_offload( + output_with_bias, forced_released_tensors=forced_released_tensors + ) + + def _forward_self_attention_output_with_bias( + self, + hidden_states: Tensor, + attention_mask: Optional[Tensor] = None, + rotary_pos_emb: Optional[Tensor] = None, + rotary_pos_cos: Optional[Tensor] = None, + rotary_pos_sin: Optional[Tensor] = None, + rotary_pos_cos_sin: Optional[Tensor] = None, + attention_bias: Optional[Tensor] = None, + inference_context: Optional[BaseInferenceContext] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + sequence_len_offset: Optional[Tensor] = None, + *, + inference_params: Optional[Any] = None, + ): + """Run input norm and self-attention, returning the raw output before BDA.""" + inference_context = deprecate_inference_params(inference_context, inference_params) + + attn_norm_manager = self.off_interface(self.offload_attn_norm, hidden_states, "attn_norm") + if self.recompute_input_layernorm: + self.input_layernorm_checkpoint = tensor_parallel.CheckpointWithoutOutput() + with attn_norm_manager as hidden_states: + input_layernorm_output = self.input_layernorm_checkpoint.checkpoint( + apply_module(self.input_layernorm), hidden_states + ) + else: + with attn_norm_manager as hidden_states: + input_layernorm_output = apply_module(self.input_layernorm)(hidden_states) + + if isinstance(input_layernorm_output, tuple): + if len(input_layernorm_output) != 2: + raise ValueError( + "When the output of input_layernorm is a tuple, it is expected " + f"to have 2 elements (output, residual), but got " + f"{len(input_layernorm_output)}" + ) + input_layernorm_output, residual = input_layernorm_output + else: + residual = hidden_states + + if self.config.fp32_residual_connection: + residual = residual.float() + + using_fused_tp_inference_kernel = ( + InferenceMode.is_active() and self.config.inference_fuse_tp_communication + ) + if using_fused_tp_inference_kernel: + self._set_proj_residual(residual) + + nvtx_range_push(suffix="self_attention") + attention_output_with_bias = self.self_attention( + input_layernorm_output, + attention_mask=attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + rotary_pos_cos_sin=rotary_pos_cos_sin, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + sequence_len_offset=sequence_len_offset, + ) + nvtx_range_pop(suffix="self_attention") + + if self.recompute_input_layernorm: + self.input_layernorm_checkpoint.discard_output_and_register_recompute( + attention_output_with_bias[0] + ) + + return attention_output_with_bias, attn_norm_manager, residual + def _forward_attention( self, hidden_states: Tensor, @@ -581,6 +685,7 @@ def _forward_attention( packed_seq_params: Optional[PackedSeqParams] = None, sequence_len_offset: Optional[Tensor] = None, padding_mask: Optional[Tensor] = None, + input_ids: Optional[Tensor] = None, *, inference_params: Optional[Any] = None, ): @@ -604,6 +709,7 @@ def _forward_attention( packed_seq_params (object, optional): Parameters for packed sequence processing. sequence_len_offset (Tensor, optional): Offset along sequence dimension during inference. + input_ids (Tensor, optional): Token IDs forwarded to hash-routed MoE layers. Returns: Tuple[Tensor, Tensor]: A tuple containing: @@ -612,43 +718,16 @@ def _forward_attention( otherwise None. """ inference_context = deprecate_inference_params(inference_context, inference_params) - - # Optional Input Layer norm - attn_norm_manager = self.off_interface(self.offload_attn_norm, hidden_states, "attn_norm") - if self.recompute_input_layernorm: - self.input_layernorm_checkpoint = tensor_parallel.CheckpointWithoutOutput() - with attn_norm_manager as hidden_states: - input_layernorm_output = self.input_layernorm_checkpoint.checkpoint( - apply_module(self.input_layernorm), hidden_states - ) - else: - with attn_norm_manager as hidden_states: - input_layernorm_output = apply_module(self.input_layernorm)(hidden_states) - - if isinstance(input_layernorm_output, tuple): - if len(input_layernorm_output) != 2: - raise ValueError( - f"When the output of input_layernorm is a tuple, it is " - f"expected to have 2 elements (output, residual), but " - f"got {len(input_layernorm_output)}" - ) - input_layernorm_output, residual = input_layernorm_output - else: - residual = hidden_states - - if self.config.fp32_residual_connection: - residual = residual.float() + input_layernorm_output, residual, attn_state = self._run_input_layernorm(hidden_states) using_fused_tp_inference_kernel = ( InferenceMode.is_active() and self.config.inference_fuse_tp_communication ) - if using_fused_tp_inference_kernel: # Set the residual for fused reduce-scatter + add + layer-norm + all-gather # operation in attention's out_proj (linear_proj) self._set_proj_residual(residual) - # Self attention. nvtx_range_push(suffix="self_attention") attention_output_with_bias = self.self_attention( input_layernorm_output, @@ -664,13 +743,71 @@ def _forward_attention( ) nvtx_range_pop(suffix="self_attention") - if self.recompute_input_layernorm: + if self._input_layernorm_checkpoint_active: # discard the output of the input layernorm and register the recompute # as a gradient hook of attention_output_with_bias[0] self.input_layernorm_checkpoint.discard_output_and_register_recompute( attention_output_with_bias[0] ) + hidden_states = self._apply_self_attn_bda_step( + attention_output_with_bias, residual, attn_state + ) + return self._run_cross_attention(hidden_states, context, context_mask, inference_context) + + def _run_input_layernorm(self, hidden_states): + """Run input layernorm with optional output-discarding checkpoint and + fine-grained activation offloading. + + Sets ``self._input_layernorm_checkpoint_active`` so the caller can gate + the post-attention discard-and-register hook on the same condition. The + flag is consumed by the next ``self._apply_self_attn_bda_step`` step. + + Returns: + Tuple ``(input_layernorm_output, residual, attn_state)`` where + ``attn_state`` is an opaque payload subclasses can use to thread + extra intermediates (e.g. mHC ``h_res``/``h_post``) through to + ``_apply_self_attn_bda_step``. Base returns ``()``. + """ + self.attn_norm_manager = self.off_interface( + self.offload_attn_norm, hidden_states, "attn_norm" + ) + self._input_layernorm_checkpoint_active = self.recompute_input_layernorm + if self._input_layernorm_checkpoint_active: + self.input_layernorm_checkpoint = tensor_parallel.CheckpointWithoutOutput() + with self.attn_norm_manager as hidden_states: + input_layernorm_output = self.input_layernorm_checkpoint.checkpoint( + apply_module(self.input_layernorm), hidden_states + ) + else: + with self.attn_norm_manager as hidden_states: + input_layernorm_output = apply_module(self.input_layernorm)(hidden_states) + + if isinstance(input_layernorm_output, tuple): + if len(input_layernorm_output) != 2: + raise ValueError( + f"When the output of input_layernorm is a tuple, it is " + f"expected to have 2 elements (output, residual), but " + f"got {len(input_layernorm_output)}" + ) + input_layernorm_output, residual = input_layernorm_output + else: + residual = hidden_states + + if self.config.fp32_residual_connection: + residual = residual.float() + return input_layernorm_output, residual, () + + def _apply_self_attn_bda_step(self, attention_output_with_bias, residual, attn_state=()): + """bias-dropout-add for self-attention output + post-step offload commit. + + Subclasses override this to swap in a fused kernel that consumes extra + intermediates threaded via ``attn_state`` (the third element returned + by ``_run_input_layernorm``). Base ignores ``attn_state``. + """ + using_fused_tp_inference_kernel = ( + InferenceMode.is_active() and self.config.inference_fuse_tp_communication + ) # TODO: could we move `bias_dropout_add_exec_handler` itself # inside the module provided in the `bias_dropout_add_spec` module? nvtx_range_push(suffix="self_attn_bda") @@ -688,11 +825,14 @@ def _forward_attention( # Delay the offload of the attention norm until after the self_attn_bda has been computed # because the residual is needed in the self_attn_bda. - hidden_states = attn_norm_manager.group_offload( + hidden_states = self.attn_norm_manager.group_offload( hidden_states, forced_released_tensors=[residual] ) + self.attn_norm_manager = None + return hidden_states - # Optional Layer norm after self-attention + def _run_cross_attention(self, hidden_states, context, context_mask, inference_context): + """Optional pre-cross-attn layernorm + cross-attention + bda block.""" pre_cross_attn_layernorm_output = apply_module(self.pre_cross_attn_layernorm)(hidden_states) if isinstance(pre_cross_attn_layernorm_output, tuple): @@ -709,7 +849,6 @@ def _forward_attention( if self.config.fp32_residual_connection: residual = residual.float() - # Cross attention. attention_output_with_bias = self.cross_attention( pre_cross_attn_layernorm_output, attention_mask=context_mask, @@ -737,12 +876,20 @@ def forward(self, *args, **kwargs): This method calls the core computation of a transformer layer, including self-attention, cross-attention (if applicable), and feed-forward operations. """ + called_from_hybrid_mhc_wrapper = kwargs.pop("_called_from_hybrid_mhc_wrapper", False) + if self.config.enable_hyper_connections and not called_from_hybrid_mhc_wrapper: + raise RuntimeError( + "TransformerLayer.forward() must not be called directly when " + "enable_hyper_connections=True. HyperConnectionHybridLayer must drive " + "the wrapped TransformerLayer through this path." + ) hidden_states, context = self._forward_attention(*args, **kwargs) output = self._forward_mlp( hidden_states, kwargs.get("inference_context", None), padding_mask=kwargs.get("padding_mask", None), packed_seq_params=kwargs.get("packed_seq_params", None), + input_ids=kwargs.get("input_ids", None), ) return output, context @@ -799,12 +946,90 @@ def _maybe_reflatten_from_moe(self, output, packed_seq_params, mbs): return output return output.transpose(0, 1).reshape(mbs * packed_seq_params.tokens_per_sample, 1, -1) + def _run_pre_mlp_layernorm(self, hidden_states): + """Run pre-MLP layernorm (with optional recompute and offload), unpack a + tuple-output layernorm, and apply the fp32-residual cast. + + Returns: + Tuple ``(pre_mlp_layernorm_output, residual, mlp_state)`` where + ``mlp_state`` is an opaque payload subclasses can use to thread + extra intermediates (e.g. mHC ``mlp_h_res`` / ``mlp_hc_h_post``) + through to ``_apply_mlp_bda_step``. Base returns ``()``. + """ + pre_mlp_layernorm_output = self._forward_pre_mlp_layernorm(hidden_states) + + if isinstance(pre_mlp_layernorm_output, tuple): + if len(pre_mlp_layernorm_output) != 2: + raise ValueError( + f"When the output of pre_mlp_layernorm is a tuple, it is " + f"expected to have 2 elements (output, residual), but " + f"got {len(pre_mlp_layernorm_output)}" + ) + pre_mlp_layernorm_output, residual = pre_mlp_layernorm_output + else: + # Residual connection. + residual = hidden_states + + if self.config.fp32_residual_connection: + residual = residual.float() + + return pre_mlp_layernorm_output, residual, () + + def _forward_mlp_output_with_bias( + self, + hidden_states: Tensor, + inference_context: BaseInferenceContext | None = None, + padding_mask: Tensor | None = None, + packed_seq_params=None, + input_ids: Optional[Tensor] = None, + ) -> tuple[tuple[Tensor, Tensor | None], Tensor]: + """Run pre-MLP norm and MLP/MoE, returning the raw output before BDA.""" + pre_mlp_layernorm_output = self._forward_pre_mlp_layernorm(hidden_states) + + if isinstance(pre_mlp_layernorm_output, tuple): + if len(pre_mlp_layernorm_output) != 2: + raise ValueError( + "When the output of pre_mlp_layernorm is a tuple, it is expected " + f"to have 2 elements (output, residual), but got " + f"{len(pre_mlp_layernorm_output)}" + ) + pre_mlp_layernorm_output, residual = pre_mlp_layernorm_output + else: + residual = hidden_states + + if self.config.fp32_residual_connection: + residual = residual.float() + + pre_mlp_layernorm_output, padding_mask, moe_unflatten_mbs = ( + self._maybe_unflatten_for_moe( + pre_mlp_layernorm_output, padding_mask, packed_seq_params + ) + ) + + mlp_output_with_bias = self._run_mlp( + pre_mlp_layernorm_output, + residual, + padding_mask, + inference_context, + input_ids=input_ids, + ) + + if moe_unflatten_mbs is not None: + mlp_output, mlp_bias = mlp_output_with_bias + mlp_output = self._maybe_reflatten_from_moe( + mlp_output, packed_seq_params, moe_unflatten_mbs + ) + mlp_output_with_bias = (mlp_output, mlp_bias) + + return mlp_output_with_bias, residual + def _forward_mlp( self, hidden_states: Tensor, inference_context: BaseInferenceContext | None = None, padding_mask: Tensor | None = None, packed_seq_params=None, + input_ids: Optional[Tensor] = None, ) -> Tensor | list[Tensor | None]: """ Perform a forward pass through the feed-forward layer. @@ -819,32 +1044,68 @@ def _forward_mlp( The MoELayer will internally transform this to [seq_length, bsz] format. packed_seq_params: Packed sequence parameters, used to detect flattened batches that need reshaping for MoE sequence load balancing. + input_ids (Tensor, optional): Token IDs with shape [batch_size, seq_length]. + Required by hash-routed MoE layers. Returns: output (Tensor): Transformed hidden states of shape [s, b, h]. """ + pre_mlp_layernorm_output, residual, mlp_state = self._run_pre_mlp_layernorm(hidden_states) - # Optional Layer norm post the cross-attention. - pre_mlp_layernorm_output = self._forward_pre_mlp_layernorm(hidden_states) + pre_mlp_layernorm_output, padding_mask, moe_unflatten_mbs = self._maybe_unflatten_for_moe( + pre_mlp_layernorm_output, padding_mask, packed_seq_params + ) - if isinstance(pre_mlp_layernorm_output, tuple): - if len(pre_mlp_layernorm_output) != 2: - raise ValueError( - f"When the output of pre_mlp_layernorm is a tuple, it is " - f"expected to have 2 elements (output, residual), but " - f"got {len(pre_mlp_layernorm_output)}" - ) - pre_mlp_layernorm_output, residual = pre_mlp_layernorm_output + mlp_output_with_bias = self._run_mlp( + pre_mlp_layernorm_output, + residual, + padding_mask, + inference_context, + input_ids=input_ids, + ) + + if moe_unflatten_mbs is not None: + mlp_output, mlp_bias = mlp_output_with_bias + mlp_output = self._maybe_reflatten_from_moe( + mlp_output, packed_seq_params, moe_unflatten_mbs + ) + mlp_output_with_bias = (mlp_output, mlp_bias) + + if ( + self.is_moe_layer + and self.config.cuda_graph_impl == "transformer_engine" + and self.training + and is_graph_capturing() + and CudaGraphModule.moe_router in self.config.cuda_graph_modules + ): + if self.recompute_pre_mlp_layernorm: + # Register the recompute hooks to all the cudagraph output tensors, because some + # tensors are in parallel execution paths and they all need pre_mlp_layernorm to be + # recomputed in backward pass. For example, the router path and the shared expert + # path. So only register in one path is risky. + for tensor in mlp_output_with_bias: + self.pre_mlp_norm_checkpoint.discard_output_and_register_recompute(tensor) + return list(mlp_output_with_bias) + [residual] else: - # Residual connection. - residual = hidden_states + return self._apply_mlp_bda_step(mlp_output_with_bias, residual, mlp_state) - if self.config.fp32_residual_connection: - residual = residual.float() + def _run_mlp( + self, + pre_mlp_layernorm_output: Tensor, + residual: Tensor, + padding_mask: Tensor | None, + inference_context: BaseInferenceContext | None, + input_ids: Optional[Tensor] = None, + ): + """Execute the MLP submodule with the appropriate variant. - pre_mlp_layernorm_output, padding_mask, moe_unflatten_mbs = self._maybe_unflatten_for_moe( - pre_mlp_layernorm_output, padding_mask, packed_seq_params - ) + Picks between the recompute (te_checkpoint / tensor_parallel.checkpoint), + chunked-prefill, and direct-call paths. Shared by both + :class:`TransformerLayer` and :class:`HyperConnectionTransformerLayer` so + the MLP-call branching stays in one place. + Returns: + ``mlp_output_with_bias``: tuple of (mlp_output, mlp_bias). + """ nvtx_range_push(suffix="mlp") # Potentially chunk the MLP computation during prefill to minimize the peak activation size should_chunk_mlp_for_prefill = ( @@ -865,6 +1126,10 @@ def _forward_mlp( InferenceMode.is_active() and self.config.inference_fuse_tp_communication ) + moe_kwargs = {} + if self.is_moe_layer and input_ids is not None: + moe_kwargs["input_ids"] = input_ids + if self.recompute_mlp: if self.config.fp8 or self.config.fp4: # import here to avoid circular import @@ -877,10 +1142,13 @@ def _forward_mlp( self.pg_collection.tp, pre_mlp_layernorm_output, padding_mask=padding_mask, + **moe_kwargs, ) else: mlp_output_with_bias = tensor_parallel.checkpoint( - functools.partial(apply_module(self.mlp), padding_mask=padding_mask), + functools.partial( + apply_module(self.mlp), padding_mask=padding_mask, **moe_kwargs + ), False, pre_mlp_layernorm_output, ) @@ -912,49 +1180,52 @@ def _forward_mlp( # operation in MLP's fc2. self._set_fc2_residual(residual) mlp_output_with_bias = apply_module(self.mlp)( - pre_mlp_layernorm_output, padding_mask=padding_mask + pre_mlp_layernorm_output, padding_mask=padding_mask, **moe_kwargs ) - if moe_unflatten_mbs is not None: - mlp_output, mlp_bias = mlp_output_with_bias - mlp_output = self._maybe_reflatten_from_moe( - mlp_output, packed_seq_params, moe_unflatten_mbs - ) - mlp_output_with_bias = (mlp_output, mlp_bias) - nvtx_range_pop(suffix="mlp") + return mlp_output_with_bias - if ( - self.is_moe_layer - and self.config.cuda_graph_impl == "transformer_engine" - and self.training - and is_graph_capturing() - and CudaGraphModule.moe_router in self.config.cuda_graph_modules - ): - if self.recompute_pre_mlp_layernorm: - # Register the recompute hooks to all the cudagraph output tensors, because some - # tensors are in parallel execution paths and they all need pre_mlp_layernorm to be - # recomputed in backward pass. For example, the router path and the shared expert - # path. So only register in one path is risky. - for tensor in mlp_output_with_bias: - self.pre_mlp_norm_checkpoint.discard_output_and_register_recompute(tensor) - return list(mlp_output_with_bias) + [residual] - else: - return self._forward_post_mlp(mlp_output_with_bias, residual) - - def _forward_post_mlp( - self, mlp_output_with_bias: tuple[Tensor, Tensor | None], residual: Tensor + def _apply_mlp_bda_step( + self, + mlp_output_with_bias: tuple[Tensor, Tensor | None], + residual: Tensor, + mlp_state: tuple = (), ) -> Tensor: """ - Perform operations after the MLP computation. + Perform operations after the MLP computation: bias-dropout-add for + the MLP output + post-step offload commit + viewless-tensor wrap. + + Subclasses override this to swap in a fused kernel that consumes extra + intermediates threaded via ``mlp_state`` (the third element returned + by ``_run_pre_mlp_layernorm``). Base ignores ``mlp_state``. Args: mlp_output_with_bias (Tensor): Output tensor of the MLP layer with bias. residual (Tensor): Residual tensor. + mlp_state: Opaque payload from ``_run_pre_mlp_layernorm``. Default ``()``. Returns: output (Tensor): Transformed hidden states of shape [s, b, h]. """ + # Back-compat shim: prior to the MLP-hook refactor this method was named + # `_forward_post_mlp` and took only (mlp_output_with_bias, residual). If a + # subclass still overrides the legacy name, route through it and emit a + # DeprecationWarning. `mlp_state` is dropped — the legacy contract didn't + # have it. To be removed in a future release. + for klass in type(self).__mro__: + if klass is TransformerLayer: + break + if "_forward_post_mlp" in vars(klass): + warnings.warn( + "TransformerLayer._forward_post_mlp has been renamed to " + "_apply_mlp_bda_step and gained an `mlp_state` parameter. " + "Override `_apply_mlp_bda_step` instead; the legacy hook " + "will be removed in a future release.", + DeprecationWarning, + stacklevel=2, + ) + return klass._forward_post_mlp(self, mlp_output_with_bias, residual) using_fused_tp_inference_kernel = ( InferenceMode.is_active() and self.config.inference_fuse_tp_communication @@ -1093,18 +1364,72 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): Dict[str, torch.Tensor]: A dictionary containing the static inputs for the layer. """ static_inputs = super().get_layer_static_inputs(seq_length, micro_batch_size) + device = torch.cuda.current_device() - if not isinstance(self.self_attention, IdentityOp) and ( + attn_in_graph = not isinstance(self.self_attention, IdentityOp) and ( not self.config.cuda_graph_modules or CudaGraphModule.attn in self.config.cuda_graph_modules - ): + ) + thd_local_seqlen = None + if self._is_thd_cuda_graph(): + if attn_in_graph: + max_seqlen = ( + self.config.max_seqlen_per_dp_cp_rank + * self.config.context_parallel_size + ) + max_num_seqs = self.config.thd_max_packed_sequences + cu_seqlens = torch.zeros( + max_num_seqs + 1, dtype=torch.int32, device=device + ) + cu_seqlens[1:] = max_seqlen + static_inputs["packed_seq_params"] = PackedSeqParams( + qkv_format="thd", + cp_partition_mode=self.config.cp_partition_mode, + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens.clone(), + cu_seqlens_q_padded=cu_seqlens.clone(), + cu_seqlens_kv_padded=cu_seqlens.clone(), + max_seqlen_q=max_seqlen, + max_seqlen_kv=max_seqlen, + pad_between_seqs=False, + ) + + thd_local_seqlen = self.config.max_seqlen_per_dp_cp_rank + if self.config.sequence_parallel: + thd_local_seqlen //= self.config.tensor_model_parallel_size + static_inputs["padding_mask"] = torch.zeros( + (1, thd_local_seqlen), dtype=torch.bool, device=device + ) + elif attn_in_graph and self.config.create_attention_mask_in_dataloader: slen_per_cp = seq_length // self.config.context_parallel_size static_inputs["attention_mask"] = ( ~(torch.tril(torch.ones((slen_per_cp, seq_length))).bool()) - .to(torch.cuda.current_device()) + .to(device) .reshape(1, 1, slen_per_cp, seq_length) .tile(micro_batch_size, 1, 1, 1) ) + elif attn_in_graph and self.self_attention.attn_mask_type not in ( + AttnMaskType.causal, + AttnMaskType.no_mask, + AttnMaskType.causal_bottom_right, + ): + log_single_rank( + logger, + logging.WARNING, + "TE CUDA graph capture is omitting attention_mask because " + "create_attention_mask_in_dataloader is False, but " + f"attn_mask_type={self.self_attention.attn_mask_type.name} may require " + "an explicit mask.", + ) + + if self.is_moe_layer and getattr(self.mlp.router, 'is_hash_layer', False): + if self._is_thd_cuda_graph(): + input_ids_shape = (1, thd_local_seqlen) + else: + input_ids_shape = (micro_batch_size, seq_length) + static_inputs["input_ids"] = torch.zeros( + input_ids_shape, dtype=torch.long, device=device + ) return static_inputs def _get_submodules_under_cudagraphs(self): @@ -1135,6 +1460,110 @@ def _get_submodules_under_cudagraphs(self): submodules += [self.mlp.shared_experts] return submodules + def _set_te_cuda_graph_packed_seq_params_static_metadata( + self, static_metadata, tensor_kwarg_names=None + ): + """Store non-Tensor ``PackedSeqParams`` metadata used during TE graph capture.""" + self._te_cuda_graph_packed_seq_params_static_metadata = dict(static_metadata) + self._te_cuda_graph_packed_seq_params_tensor_kwarg_names = ( + None if tensor_kwarg_names is None else tuple(sorted(tensor_kwarg_names)) + ) + + def _get_te_cuda_graph_packed_seq_params_static_metadata(self): + """Return the static ``PackedSeqParams`` metadata used for this TE graph.""" + return getattr(self, '_te_cuda_graph_packed_seq_params_static_metadata', None) + + def _validate_te_cuda_graph_packed_seq_params_static_metadata(self, static_metadata): + """Validate that replay uses the same static packed-sequence contract as capture.""" + expected_static_metadata = self._get_te_cuda_graph_packed_seq_params_static_metadata() + assert expected_static_metadata is not None, ( + "TE CUDA graph replay received packed_seq_params, but the graph was captured without " + "packed-sequence sample inputs. Recapture the graph with matching PackedSeqParams " + "static metadata." + ) + + mismatched_fields = [] + for field_name in sorted(set(expected_static_metadata) | set(static_metadata)): + expected_value = expected_static_metadata.get(field_name) + actual_value = static_metadata.get(field_name) + if expected_value is actual_value: + continue + if expected_value != actual_value: + mismatched_fields.append(field_name) + + assert not mismatched_fields, ( + "TE CUDA graph replay received PackedSeqParams with static metadata that differs " + "from capture. Recapture the graph for changed fields: " + f"{', '.join(mismatched_fields)}." + ) + + def _get_te_cuda_graph_packed_seq_params_tensor_kwarg_names(self): + """Return flattened ``PackedSeqParams`` Tensor kwargs used for this TE graph.""" + return getattr(self, '_te_cuda_graph_packed_seq_params_tensor_kwarg_names', None) + + def _validate_te_cuda_graph_packed_seq_params_tensor_kwargs(self, tensor_kwargs): + """Validate replay uses the same flattened Tensor field set as capture.""" + expected_names = self._get_te_cuda_graph_packed_seq_params_tensor_kwarg_names() + if expected_names is None: + return + + expected_names = set(expected_names) + actual_names = set(tensor_kwargs) + missing_names = sorted(expected_names - actual_names) + extra_names = sorted(actual_names - expected_names) + assert not missing_names and not extra_names, ( + "TE CUDA graph replay received PackedSeqParams with Tensor fields that differ " + "from capture. Recapture the graph for missing fields " + f"{missing_names} and extra fields {extra_names}." + ) + + def _rebuild_te_cuda_graph_packed_seq_params(self, kwargs): + """Rebuild ``PackedSeqParams`` from flattened TE graph capture kwargs.""" + if not has_packed_seq_params_cuda_graph_kwargs(kwargs): + return + + assert kwargs.get('packed_seq_params') is None, ( + "PackedSeqParams must be passed either as flattened TE CUDA graph kwargs or as " + "packed_seq_params, but not both." + ) + static_metadata = self._get_te_cuda_graph_packed_seq_params_static_metadata() + assert static_metadata is not None, ( + "Flattened PackedSeqParams Tensor fields require static metadata captured on the " + "TransformerLayer." + ) + tensor_kwargs = { + key: value + for key, value in kwargs.items() + if key.startswith(CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX) + } + self._validate_te_cuda_graph_packed_seq_params_tensor_kwargs(tensor_kwargs) + + kwargs['packed_seq_params'] = build_packed_seq_params_from_cuda_graph_kwargs( + kwargs, static_metadata + ) + + def _flatten_te_cuda_graph_packed_seq_params(self, kwargs): + """Flatten replay-time ``PackedSeqParams`` into Tensor kwargs for TE graphs.""" + packed_seq_params = kwargs.pop('packed_seq_params', None) + expected_static_metadata = self._get_te_cuda_graph_packed_seq_params_static_metadata() + if packed_seq_params is None: + assert expected_static_metadata is None, ( + "TE CUDA graph was captured with packed_seq_params, so replay must also pass " + "packed_seq_params with matching static metadata." + ) + return + + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params) + self._validate_te_cuda_graph_packed_seq_params_static_metadata(static_metadata) + self._validate_te_cuda_graph_packed_seq_params_tensor_kwargs(tensor_kwargs) + + duplicate_keys = set(kwargs) & set(tensor_kwargs) + assert not duplicate_keys, ( + "PackedSeqParams CUDA graph Tensor kwargs overlap with existing replay kwargs: " + f"{', '.join(sorted(duplicate_keys))}." + ) + kwargs.update(tensor_kwargs) + def _te_cuda_graph_capture(self, *args, **kwargs): """ CUDA Graph capture for this layer using TE interface. @@ -1155,6 +1584,8 @@ def _te_cuda_graph_capture(self, *args, **kwargs): hidden_states = kwargs.pop("hidden_states") hidden_states = self.off_interface.backward_record(hidden_states) kwargs["hidden_states"] = hidden_states + self._rebuild_te_cuda_graph_packed_seq_params(kwargs) + context = None if ( not self.config.cuda_graph_modules @@ -1178,7 +1609,12 @@ def _te_cuda_graph_capture(self, *args, **kwargs): ) ) ): - hidden_states = self._forward_mlp(hidden_states) + hidden_states = self._forward_mlp( + hidden_states, + padding_mask=kwargs.get("padding_mask", None), + packed_seq_params=kwargs.get("packed_seq_params", None), + input_ids=kwargs.get("input_ids", None), + ) if not isinstance(hidden_states, list) and not isinstance(hidden_states, tuple): cuda_graph_outputs = [hidden_states] else: @@ -1197,36 +1633,107 @@ def _te_cuda_graph_replay(self, *args, **kwargs): CUDA graph replay for this layer and microbatch `self.current_microbatch` using TE interface. TransformerEngine versions>=1.10 allow keyword arguments with CUDA graph. However, CUDA graph accepts only Tensor inputs. - Hence, `inference_context` and `packed_seq_params` are excluded from input list. + Hence, `inference_context` is excluded from input list. `packed_seq_params` is split + into Tensor graph inputs and static metadata when attention is in the graph scope. """ context = None + padding_mask = kwargs.get("padding_mask", None) + packed_seq_params = kwargs.get("packed_seq_params", None) if ( self.config.cuda_graph_modules and CudaGraphModule.attn not in self.config.cuda_graph_modules ): + input_ids = kwargs.get("input_ids", None) hidden_states, context = self._forward_attention(*args, **kwargs) args = (hidden_states,) kwargs = {} + if padding_mask is not None: + kwargs["padding_mask"] = padding_mask + if input_ids is not None: + kwargs["input_ids"] = input_ids + else: + self._flatten_te_cuda_graph_packed_seq_params(kwargs) - assert (kwargs.get('inference_context') is None) and ( - kwargs.get('packed_seq_params') is None - ), ( + assert kwargs.get('inference_context') is None, ( "CUDA graph accepts only Tensor inputs. " - "inference_context and packed_seq_params are excluded from input list. " + "inference_context is excluded from input list; packed_seq_params must be " + "flattened into Tensor kwargs with matching static metadata. " "For inference cuda graph, please use cuda_graph_impl=local instead." ) if self.config.delay_offload_until_cuda_graph: self.off_interface.enter_replay() + self._te_cuda_graph_replay_packed_seq_params = packed_seq_params + self._te_cuda_graph_replay_padding_mask = padding_mask try: return self._te_cuda_graph_replay_impl(args, kwargs, context) finally: + self._te_cuda_graph_replay_packed_seq_params = None + self._te_cuda_graph_replay_padding_mask = None if self.config.delay_offload_until_cuda_graph: self.off_interface.exit_replay() + def resume_moe_experts_after_partial_cudagraph(self, cuda_graph_output): + """Resume eager MoE expert computation from captured router outputs. + + Returns the raw ``(output, bias)`` pair so callers can choose the + appropriate residual/BDA owner. + """ + assert self.is_moe_layer + assert CudaGraphModule.moe_router in self.config.cuda_graph_modules + assert not self.config.overlap_moe_expert_parallel_comm, ( + "Hybrid mHC partial-MoE replay does not support EP overlap." + ) + + cuda_graph_output = list(cuda_graph_output) + residual = cuda_graph_output.pop() + shared_expert_output, routing_map = None, None + if ( + self.config.moe_shared_expert_intermediate_size is not None + and not self.config.moe_shared_expert_overlap + ): + shared_expert_output = cuda_graph_output.pop() + + if CudaGraphModule.moe_preprocess in self.config.cuda_graph_modules: + (hidden_states, probs), attr_outputs = ( + cuda_graph_output[:2], + cuda_graph_output[2:], + ) + valid_attrs = self.mlp.token_dispatcher.valid_cudagraph_attrs + assert len(attr_outputs) == len(valid_attrs) + for attr_name, attr_value in zip(valid_attrs, attr_outputs): + attr = self.mlp.token_dispatcher + path = attr_name.split('.') + for name in path[:-1]: + attr = getattr(attr, name) + setattr(attr, path[-1], attr_value) + else: + assert len(cuda_graph_output) == 3, ( + "Partial MoE graph output must be [hidden_states, probs, routing_map]." + ) + hidden_states, probs, routing_map = cuda_graph_output + + nvtx_range_push(suffix="mlp") + self.mlp.cudagraph_tensor_store.set( + hidden_states=hidden_states, + probs=probs, + routing_map=routing_map, + shared_expert_output=shared_expert_output, + ) + try: + mlp_output_with_bias = apply_module(self.mlp)(hidden_states) + finally: + self.mlp.cudagraph_tensor_store.clear() + nvtx_range_pop(suffix="mlp") + return residual, mlp_output_with_bias + def _te_cuda_graph_replay_impl(self, args, kwargs, context): """Implementation of _te_cuda_graph_replay, separated for replay mode cleanup.""" + packed_seq_params = getattr( + self, '_te_cuda_graph_replay_packed_seq_params', None + ) + padding_mask = getattr(self, '_te_cuda_graph_replay_padding_mask', None) cuda_graph_output = list(super()._te_cuda_graph_replay(*args, **kwargs)) # Flush delayed offload groups from previous layers after graph replay. @@ -1253,63 +1760,58 @@ def _te_cuda_graph_replay_impl(self, args, kwargs, context): elif self.is_moe_layer and CudaGraphModule.moe_router in self.config.cuda_graph_modules: # CUDA Graph partially captures the MoE. # The rest of the layer should go to the normal pass. - shared_expert_output, routing_map = None, None - # residual is the last element in the CUDA graph output. - residual = cuda_graph_output.pop() - if ( - self.config.moe_shared_expert_intermediate_size is not None - and not self.config.moe_shared_expert_overlap - ): - # The shared expert output is the last second element in the CUDA graph output. - shared_expert_output = cuda_graph_output.pop() - - if CudaGraphModule.moe_preprocess in self.config.cuda_graph_modules: - # CUDA graph output is [hidden_states, probs] + attributes outputs. - (hidden_states, probs), attr_outputs = cuda_graph_output[:2], cuda_graph_output[2:] - valid_cudagraph_attrs = self.mlp.token_dispatcher.valid_cudagraph_attrs - assert len(attr_outputs) == len( - valid_cudagraph_attrs - ), f"attr_outputs: {len(attr_outputs)} != {len(valid_cudagraph_attrs)}" - for i, attr_name in enumerate(valid_cudagraph_attrs): - hier_attr_name = attr_name.split('.') - attr = self.mlp.token_dispatcher - for name in hier_attr_name[:-1]: - attr = getattr(attr, name) - setattr(attr, hier_attr_name[-1], attr_outputs[i]) - else: - # CUDA graph output is [hidden_states, probs, routing_map]. - assert len(cuda_graph_output) == 3, ( - "CUDA graph output should be [hidden_states, probs, routing_map], " - f"but got {len(cuda_graph_output)} elements" - ) - hidden_states, probs, routing_map = cuda_graph_output - - # Resume the MoELayer forward pass from the end of the CUDA graph scope. - # The MoE layer will skip redundant computations when we pass in the calculated values - # through the keyword arguments. See MoELayer.forward docstring for more details. - nvtx_range_push(suffix="mlp") - self.mlp.cudagraph_tensor_store.set( - hidden_states=hidden_states, - probs=probs, - routing_map=routing_map, - shared_expert_output=shared_expert_output, - ) # If EP overlap is enabled, remaining of mlp will be called as fine_grained_callables # and should be skipped here. if self.config.overlap_moe_expert_parallel_comm: - probs, routing_map = self.mlp.route(hidden_states) - hidden_states, probs = self.mlp.preprocess(hidden_states, probs, routing_map) - nvtx_range_pop(suffix="mlp") + overlap_outputs = list(cuda_graph_output) + residual = overlap_outputs.pop() + shared_expert_output = None + if ( + self.config.moe_shared_expert_intermediate_size is not None + and not self.config.moe_shared_expert_overlap + ): + shared_expert_output = overlap_outputs.pop() + routing_map = None + if CudaGraphModule.moe_preprocess in self.config.cuda_graph_modules: + (hidden_states, probs), attr_outputs = ( + overlap_outputs[:2], + overlap_outputs[2:], + ) + valid_attrs = self.mlp.token_dispatcher.valid_cudagraph_attrs + assert len(attr_outputs) == len(valid_attrs) + for attr_name, attr_value in zip(valid_attrs, attr_outputs): + attr = self.mlp.token_dispatcher + path = attr_name.split('.') + for name in path[:-1]: + attr = getattr(attr, name) + setattr(attr, path[-1], attr_value) + else: + assert len(overlap_outputs) == 3 + hidden_states, probs, routing_map = overlap_outputs + self.mlp.cudagraph_tensor_store.set( + hidden_states=hidden_states, + probs=probs, + routing_map=routing_map, + shared_expert_output=shared_expert_output, + ) + nvtx_range_push(suffix="mlp") + try: + probs, routing_map = self.mlp.route(hidden_states) + hidden_states, probs = self.mlp.preprocess( + hidden_states, probs, routing_map + ) + finally: + nvtx_range_pop(suffix="mlp") return residual, hidden_states, probs, shared_expert_output - mlp_output_with_bias = apply_module(self.mlp)(hidden_states) - self.mlp.cudagraph_tensor_store.clear() - nvtx_range_pop(suffix="mlp") + residual, mlp_output_with_bias = ( + self.resume_moe_experts_after_partial_cudagraph(cuda_graph_output) + ) # If we early returned, layernorm recompute hooks were attached to the output buffer - # of the cudagraph, so disable the recompute hooks inside _forward_post_mlp + # of the cudagraph, so disable the recompute hooks inside _apply_mlp_bda_step recompute_pre_mlp_layernorm = self.recompute_pre_mlp_layernorm self.recompute_pre_mlp_layernorm = False - output = self._forward_post_mlp(mlp_output_with_bias, residual) + output = self._apply_mlp_bda_step(mlp_output_with_bias, residual) self.recompute_pre_mlp_layernorm = recompute_pre_mlp_layernorm else: # If EP overlap is enabled, needs to return same outputs as submodule.attn @@ -1334,7 +1836,13 @@ def _te_cuda_graph_replay_impl(self, args, kwargs, context): return residual, hidden_states, probs, shared_expert_output # CUDA Graph does not capture the MLP/MoE part at all. - output = self._forward_mlp(*cuda_graph_output) + assert len(cuda_graph_output) >= 1 + output = self._forward_mlp( + *cuda_graph_output, + padding_mask=padding_mask, + packed_seq_params=packed_seq_params, + input_ids=kwargs.get("input_ids", None), + ) return output, context def _get_te_cuda_graph_replay_args(self, *args, **kwargs): @@ -1383,11 +1891,12 @@ def get_zero_attention_mask(slen_per_tpcp, micro_batch_size): 'attention_mask' in cudagraph_kwargs and cudagraph_kwargs['attention_mask'] is None ): # The attention_mask can be None when there is no padding to the input sequence. - # However, an attention_mask Tensor must be passed into cudagraph for replay, so - # we create an equivalent zero Tensor as the attention_mask. - cudagraph_kwargs["attention_mask"] = get_zero_attention_mask( - hidden_states.size(0), hidden_states.size(1) - ) + if not self.config.create_attention_mask_in_dataloader: + cudagraph_kwargs.pop("attention_mask") + else: + cudagraph_kwargs["attention_mask"] = get_zero_attention_mask( + hidden_states.size(0), hidden_states.size(1) + ) except ImportError: raise RuntimeError("CUDAGraph requires TransformerEngine, but not installed") return tuple(cudagraph_args), cudagraph_kwargs @@ -1508,6 +2017,316 @@ def get_layer_norm_weights(self): return +class HyperConnectionTransformerLayer(TransformerLayer): + """A transformer layer with Manifold-Constrained Hyper-Connections (mHC). + + Extends TransformerLayer by adding hyper connection modules around self-attention + and MLP. The n-stream hidden states are aggregated before each sub-layer and + expanded back afterwards using learned mappings (H_pre, H_post, H_res). + + Cross-attention hyper connection is not supported. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: TransformerLayerSubmodules, + layer_number: int = 1, + hidden_dropout: Optional[float] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + vp_stage: Optional[int] = None, + ): + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + hidden_dropout=hidden_dropout, + pg_collection=pg_collection, + vp_stage=vp_stage, + ) + + if submodules.cross_attention_hyper_connection is not IdentityOp: + raise ValueError( + "HyperConnectionTransformerLayer does not support cross-attention " + "hyper connections. Use IdentityOp for cross_attention_hyper_connection." + ) + + assert submodules.self_attention_hyper_connection is not IdentityOp, ( + "HyperConnectionTransformerLayer requires self_attention_hyper_connection. " + "Use TransformerLayer instead if hyper connections are not needed." + ) + assert submodules.mlp_hyper_connection is not IdentityOp, ( + "HyperConnectionTransformerLayer requires mlp_hyper_connection. " + "Use TransformerLayer instead if hyper connections are not needed." + ) + + # mHC over a single MoE-MLP layer is not supported in this implementation; + # compose mHC with MoE by wrapping MoE inside a HyperConnectionHybridLayer + # (HybridStack path) instead. This guard fires at setup so misconfigured + # specs fail fast rather than producing silently-wrong shapes at runtime. + if self.is_moe_layer: + raise NotImplementedError( + "HyperConnectionTransformerLayer does not support MoE MLP submodules. " + "To combine mHC with MoE, wrap the MoE block as a HybridStack layer " + "via HyperConnectionHybridLayer instead." + ) + + self.self_attention_hyper_connection = build_module( + submodules.self_attention_hyper_connection, + config=self.config, + layer_number=self.layer_number, + ) + + self.mlp_hyper_connection = build_module( + submodules.mlp_hyper_connection, config=self.config, layer_number=self.layer_number + ) + + # When mHC recompute is active, skip checkpointing if the layernorm + # is IdentityOp (fused into TE linear) — there is nothing to recompute. + self.mhc_checkpoint_input_layernorm = not isinstance(self.input_layernorm, IdentityOp) + self.mhc_checkpoint_pre_mlp_layernorm = not isinstance(self.pre_mlp_layernorm, IdentityOp) + + # Set per-call by __call__ from kwargs so forward can read it without re-piping + # the manager through the CUDA-graph kwarg path (CheckpointWithoutOutputManager + # is not a CUDA-graph-supported type and gets stripped during capture). Read by + # _run_input_layernorm, _apply_self_attn_bda_step, _run_pre_mlp_layernorm, and + # _apply_mlp_bda_step — do not delete; appears unused only at the class level. + self._mhc_recompute_manager: Optional['CheckpointWithoutOutputManager'] = None + + def __call__(self, *args, **kwargs): + # Pull the manager off kwargs before super().__call__ hands them to the + # CUDA-graph machinery (which can't handle a CheckpointWithoutOutputManager). + # forward() reads the value back from self. + self._mhc_recompute_manager = kwargs.pop("mhc_recompute_manager", None) + return super().__call__(*args, **kwargs) + + def get_layer_static_inputs(self, seq_length, micro_batch_size): + """Override to produce n-stream hidden_states of shape [s, b, n*C]. + + CUDA graph capture creates static buffers whose shapes are determined by + this method. The base class returns [s, b, C], but mHC layers operate on + n-stream hidden states of shape [s, b, n*C]. + """ + static_inputs = super().get_layer_static_inputs(seq_length, micro_batch_size) + hs = static_inputs["hidden_states"] + n = self.config.num_residual_streams + static_inputs["hidden_states"] = torch.ones( + (hs.shape[0], hs.shape[1], n * self.config.hidden_size), + dtype=hs.dtype, + requires_grad=hs.requires_grad, + device=hs.device, + ) + return static_inputs + + def _get_submodules_under_cudagraphs(self): + """Override to include hyper connection modules. + + The base TransformerLayer._get_submodules_under_cudagraphs does not include + self_attention_hyper_connection / mlp_hyper_connection. Their learnable + parameters (mapping_proj, alpha_*, bias) need manual pre-forward hooks + during CUDA graph replay so that parameter all-gathers are triggered. + """ + submodules = super()._get_submodules_under_cudagraphs() + + if not self.config.cuda_graph_modules: + return submodules + + if CudaGraphModule.attn in self.config.cuda_graph_modules: + submodules.append(self.self_attention_hyper_connection) + # HC layer rejects MoE MLPs in __init__, so only the dense (mlp) scope applies. + if CudaGraphModule.mlp in self.config.cuda_graph_modules: + submodules.append(self.mlp_hyper_connection) + return submodules + + def forward(self, *args, **kwargs): + """Forward pass with MHC recompute manager support. + + Inherits ``_forward_attention`` and ``_forward_mlp`` from base; the + mHC-specific behavior is contained in the ``_run_input_layernorm``, + ``_apply_self_attn_bda_step``, ``_run_pre_mlp_layernorm``, and + ``_apply_mlp_bda_step`` overrides, which read the manager off + ``self`` and thread per-call intermediates through the + ``attn_state`` / ``mlp_state`` slots. + + Override exists only to skip the ``enable_hyper_connections`` assert + on base ``TransformerLayer.forward``. + """ + hidden_states, context = self._forward_attention(*args, **kwargs) + output = self._forward_mlp( + hidden_states, + kwargs.get("inference_context", None), + padding_mask=kwargs.get("padding_mask", None), + packed_seq_params=kwargs.get("packed_seq_params", None), + ) + return output, context + + def _run_input_layernorm(self, hidden_states): + """HC input layernorm: hyper-connection pre-wrap + mHC-aware checkpoint. + + Threads ``h_res`` and ``h_post`` (produced by the hyper-connection + pre-wrap) to ``_apply_self_attn_bda_step`` via the ``attn_state`` slot + in the return tuple. Also sets + ``self._input_layernorm_checkpoint_active`` for the post-self-attn + discard hook. + + Returns ``(input_layernorm_output, residual, (h_res, h_post))`` where + ``residual`` is the n-stream hidden state captured before + aggregation — it flows to ``_apply_self_attn_bda_step`` via the base + skeleton's ``residual`` argument, and ``(h_res, h_post)`` flows via + ``attn_state``. + """ + # Capture the n-stream residual BEFORE self_attention_hyper_connection + # aggregates n-stream -> single-stream. The fused bda kernel needs the + # original n-stream tensor. + residual = hidden_states + + nvtx_range_push(suffix="self_attention_hyper_connection") + hidden_states, h_res, h_post = self.self_attention_hyper_connection( + hidden_states, mhc_recompute_manager=self._mhc_recompute_manager + ) + nvtx_range_pop(suffix="self_attention_hyper_connection") + + self.attn_norm_manager = self.off_interface( + self.offload_attn_norm, hidden_states, "attn_norm" + ) + self._input_layernorm_checkpoint_active = self.recompute_input_layernorm or ( + self._mhc_recompute_manager is not None and self.mhc_checkpoint_input_layernorm + ) + if self._input_layernorm_checkpoint_active: + self.input_layernorm_checkpoint = tensor_parallel.CheckpointWithoutOutput( + ckpt_manager=self._mhc_recompute_manager + ) + with self.attn_norm_manager as hidden_states: + input_layernorm_output = self.input_layernorm_checkpoint.checkpoint( + self.input_layernorm, hidden_states + ) + else: + with self.attn_norm_manager as hidden_states: + input_layernorm_output = self.input_layernorm(hidden_states) + + return input_layernorm_output, residual, (h_res, h_post) + + def _apply_self_attn_bda_step(self, attention_output_with_bias, residual, attn_state): + """HC fused bias-dropout-add: combines apply_h_res + apply_h_post + bda. + + Unpacks ``h_res`` and ``h_post`` from ``attn_state`` (threaded by + ``_run_input_layernorm`` via the base skeleton). + """ + h_res, h_post = attn_state + nvtx_range_push(suffix="self_attention_fused_h_res_h_post_bda") + with self.bias_dropout_add_exec_handler(): + hidden_states = self.self_attention_hyper_connection.fused_h_res_h_post_bda( + h_res, + residual, + h_post, + attention_output_with_bias, + self.hidden_dropout, + self.training, + self.config.bias_dropout_fusion, + self._mhc_recompute_manager, + ) + nvtx_range_pop(suffix="self_attention_fused_h_res_h_post_bda") + # HC omits forced_released_tensors — the n-stream residual is consumed + # by the fused kernel above, so the base class's "release residual after + # commit" trick doesn't apply. + hidden_states = self.attn_norm_manager.group_offload(hidden_states) + self.attn_norm_manager = None + return hidden_states + + def _run_pre_mlp_layernorm(self, hidden_states): + """HC pre-mlp layernorm: hyper-connection pre-wrap + mHC-aware checkpoint. + + Threads ``mlp_h_res`` and ``mlp_hc_h_post`` (produced by the + hyper-connection pre-wrap) to ``_apply_mlp_bda_step`` via the + ``mlp_state`` slot in the return tuple. + + Returns ``(pre_mlp_layernorm_output, residual, (mlp_h_res, mlp_hc_h_post))`` + where ``residual`` is the n-stream hidden state captured before + aggregation — it flows to ``_apply_mlp_bda_step`` via the base + skeleton's ``residual`` argument, and ``(mlp_h_res, mlp_hc_h_post)`` + flows via ``mlp_state``. + """ + # Capture the n-stream residual BEFORE mlp_hyper_connection + # aggregates n-stream -> single-stream. The fused bda kernel needs the + # original n-stream tensor. + residual = hidden_states + + nvtx_range_push(suffix="mlp_hyper_connection") + hidden_states, mlp_h_res, mlp_hc_h_post = self.mlp_hyper_connection( + hidden_states, mhc_recompute_manager=self._mhc_recompute_manager + ) + nvtx_range_pop(suffix="mlp_hyper_connection") + + self.mlp_norm_manager = self.off_interface(self.offload_mlp_norm, hidden_states, "mlp_norm") + checkpoint_pre_mlp_layernorm = self.recompute_pre_mlp_layernorm or ( + self._mhc_recompute_manager is not None and self.mhc_checkpoint_pre_mlp_layernorm + ) + if checkpoint_pre_mlp_layernorm: + self.pre_mlp_norm_checkpoint = tensor_parallel.CheckpointWithoutOutput( + ckpt_manager=self._mhc_recompute_manager + ) + with self.mlp_norm_manager as hidden_states: + pre_mlp_layernorm_output = self.pre_mlp_norm_checkpoint.checkpoint( + self.pre_mlp_layernorm, hidden_states + ) + else: + with self.mlp_norm_manager as hidden_states: + pre_mlp_layernorm_output = self.pre_mlp_layernorm(hidden_states) + + return pre_mlp_layernorm_output, residual, (mlp_h_res, mlp_hc_h_post) + + def _apply_mlp_bda_step(self, mlp_output_with_bias, residual, mlp_state): + """HC fused bias-dropout-add for MLP: combines apply_h_res + apply_h_post + bda. + + Unpacks ``mlp_h_res`` and ``mlp_hc_h_post`` from ``mlp_state`` (threaded + by ``_run_pre_mlp_layernorm`` via the base skeleton). Computes the + per-call ``mhc_mlp_bda_manager`` from ``self._mhc_recompute_manager``: + the last layer of a recompute block does NOT pass the manager into the + fused-bda checkpoint — the block-end finalize hook handles its output + discard. + """ + mlp_h_res, mlp_hc_h_post = mlp_state + + is_last_in_recompute_block = bool( + self._mhc_recompute_manager is not None + and getattr(self._mhc_recompute_manager, "is_last_layer_in_recompute_block", False) + ) + mhc_mlp_bda_manager = None if is_last_in_recompute_block else self._mhc_recompute_manager + + if self.recompute_pre_mlp_layernorm or ( + mhc_mlp_bda_manager is not None and self.mhc_checkpoint_pre_mlp_layernorm + ): + self.pre_mlp_norm_checkpoint.discard_output_and_register_recompute( + mlp_output_with_bias[0] + ) + + nvtx_range_push(suffix="mlp_fused_h_res_h_post_bda") + with self.bias_dropout_add_exec_handler(): + hidden_states = self.mlp_hyper_connection.fused_h_res_h_post_bda( + mlp_h_res, + residual, + mlp_hc_h_post, + mlp_output_with_bias, + self.hidden_dropout, + self.training, + self.config.bias_dropout_fusion, + mhc_mlp_bda_manager, + ) + nvtx_range_pop(suffix="mlp_fused_h_res_h_post_bda") + + # HC omits forced_released_tensors — the n-stream residual is consumed + # by the fused kernel above, so the base class's "release residual after + # commit" trick doesn't apply. + if self.mlp_norm_manager is not None: + hidden_states = self.mlp_norm_manager.group_offload(hidden_states) + self.mlp_norm_manager = None + + output = make_viewless_tensor( + inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True + ) + return output + + class MoETransformerLayer(TransformerLayer): """ A Transformer layer specialized for Mixture-of-Experts (MoE) architectures. @@ -1619,7 +2438,7 @@ def _restore_token_dispatcher_attrs(self): obj, name = self._resolve_token_dispatcher_attr(attr_name) setattr(obj, name, attr) - def _forward_mlp_router(self, hidden_states, padding_mask=None): + def _forward_mlp_router(self, hidden_states, padding_mask=None, input_ids=None): """ Executes the router phase of the MoE block. @@ -1644,7 +2463,10 @@ def _forward_mlp_router(self, hidden_states, padding_mask=None): residual = residual.float() router_outputs = apply_module(self.mlp)( - pre_mlp_layernorm_output, intermediate_tensors=(), padding_mask=padding_mask + pre_mlp_layernorm_output, + intermediate_tensors=(), + padding_mask=padding_mask, + input_ids=input_ids, ) if is_graph_capturing() and not is_graph_warmup(): @@ -1688,7 +2510,7 @@ def _forward_mlp_postprocess(self, residual, output, shared_expert_output, mlp_b self.mlp.fwd_execution_map = "postprocess" output = apply_module(self.mlp)(None, intermediate_tensors=(output, shared_expert_output)) - out = self._forward_post_mlp((output, mlp_bias), residual) + out = self._apply_mlp_bda_step((output, mlp_bias), residual) if is_graph_capturing() and not is_graph_warmup(): for attr_name, attr in self.token_dispatcher_attrs.items(): @@ -1700,7 +2522,12 @@ def _forward_mlp_postprocess(self, residual, output, shared_expert_output, mlp_b return out def _forward_mlp( - self, hidden_states, inference_context=None, padding_mask=None, packed_seq_params=None + self, + hidden_states, + inference_context=None, + padding_mask=None, + packed_seq_params=None, + input_ids=None, ): """ Orchestrates the MLP forward pass, handling partial CUDA graph execution logic. @@ -1717,10 +2544,10 @@ def _forward_mlp( ) def _forward_mlp_partial_cudagraphs( - hidden_states, inference_context=None, padding_mask=None + hidden_states, inference_context=None, padding_mask=None, input_ids=None ): residual, hidden_states, probs, shared_expert_output = self._forward_mlp_router( - hidden_states, padding_mask=padding_mask + hidden_states, padding_mask=padding_mask, input_ids=input_ids ) # After the router graph replays, the captured .copy_() operations that update @@ -1748,25 +2575,33 @@ def _forward_mlp_partial_cudagraphs( _forward_mlp_partial_cudagraphs, False, tensor_parallel.random.get_cuda_rng_tracker, - parallel_state.get_tensor_model_parallel_group(), + self.pg_collection.tp, hidden_states, padding_mask=padding_mask, + input_ids=input_ids, ) else: result = tensor_parallel.checkpoint( functools.partial( - _forward_mlp_partial_cudagraphs, padding_mask=padding_mask + _forward_mlp_partial_cudagraphs, + padding_mask=padding_mask, + input_ids=input_ids, ), False, hidden_states, ) else: - result = _forward_mlp_partial_cudagraphs(hidden_states, padding_mask=padding_mask) + result = _forward_mlp_partial_cudagraphs( + hidden_states, padding_mask=padding_mask, input_ids=input_ids + ) result = self._maybe_reflatten_from_moe(result, packed_seq_params, moe_unflatten_mbs) return result else: return super()._forward_mlp( - hidden_states, padding_mask=padding_mask, packed_seq_params=packed_seq_params + hidden_states, + padding_mask=padding_mask, + packed_seq_params=packed_seq_params, + input_ids=input_ids, ) diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 76791720e9d..72967b0ebf3 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -746,6 +746,34 @@ def selective_log_softmax(logits, index): return per_token_logps +def get_rl_packed_seq_params_for_cuda_graph( + seq_length: int, + device: torch.device, + sequence_packing: bool = False, + max_sequences_per_bin: int = None, +) -> PackedSeqParams: + """Build RL ``PackedSeqParams`` used to keep CUDA graph signatures stable.""" + if sequence_packing: + assert max_sequences_per_bin is not None, ( + "max_sequences_per_bin is required when sequence_packing is enabled." + ) + return get_default_packed_seq_params( + seq_length=seq_length, + max_sequences_per_bin=max_sequences_per_bin, + device=device, + ) + + cu_seqlens = torch.tensor([0, seq_length], dtype=torch.int32, device=device) + return PackedSeqParams( + qkv_format='thd', + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=seq_length, + max_seqlen_kv=seq_length, + total_tokens=seq_length, + ) + + def get_logprobs(model, tokens, position_ids, no_grad=False, sequence_packing=False, packed_seq_params=None): """Get sequence logprobs from their token ids. @@ -773,22 +801,12 @@ def get_logprobs(model, tokens, position_ids, no_grad=False, sequence_packing=Fa # graph signature matches the training forward_step in train_rl.py. # This is necessary because reference logprobs steps will reuse the training forward graph. if packed_seq_params is None: - if sequence_packing: - packed_seq_params = get_default_packed_seq_params( - seq_length=tokens.shape[1], - max_sequences_per_bin=args.rl_sequence_packing_max_sequences_per_bin, - device=tokens.device, - ) - else: - cu_seqlens = torch.tensor([0, tokens.shape[1]], dtype=torch.int32, device=tokens.device) - packed_seq_params = PackedSeqParams( - qkv_format='thd', - cu_seqlens_q=cu_seqlens, - cu_seqlens_kv=cu_seqlens, - max_seqlen_q=tokens.shape[1], - max_seqlen_kv=tokens.shape[1], - total_tokens=tokens.shape[1], - ) + packed_seq_params = get_rl_packed_seq_params_for_cuda_graph( + seq_length=tokens.shape[1], + device=tokens.device, + sequence_packing=sequence_packing, + max_sequences_per_bin=args.rl_sequence_packing_max_sequences_per_bin, + ) nvtx_range = get_nvtx_range() diff --git a/megatron/rl/sequence_packing_utils.py b/megatron/rl/sequence_packing_utils.py index ff98b0a58e2..45753de78eb 100644 --- a/megatron/rl/sequence_packing_utils.py +++ b/megatron/rl/sequence_packing_utils.py @@ -1,22 +1,24 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -import torch +import logging import math +import typing +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Tuple + import numpy as np -from typing import List, Dict, Any, Tuple, Optional +import torch from torch.utils.data import DataLoader, TensorDataset -from dataclasses import dataclass, field -from megatron.core.utils import log_single_rank -from megatron.training.global_vars import get_args, get_tokenizer -from megatron.training.utils import get_nvtx_range -from megatron.core.packed_seq_params import PackedSeqParams + from megatron.core import mpu -import logging -import typing from megatron.core.num_microbatches_calculator import ( - get_num_microbatches, - reconfigure_num_microbatches_calculator, - ) + get_num_microbatches, + reconfigure_num_microbatches_calculator, +) +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.utils import log_single_rank +from megatron.training.global_vars import get_tokenizer +from megatron.training.utils import get_nvtx_range logger = logging.getLogger(__name__) @@ -391,8 +393,6 @@ def get_default_packed_seq_params(seq_length: int, max_sequences_per_bin: int, d PackedSeqParams configured as a single unpacked sequence. """ - args = get_args() - # Pad to the maximum number of sequences in the bin for the attention kernel. # We add 2 to account for the initial 0 and the final bin_size. cu_seqlens = torch.full( diff --git a/megatron/training/argument_utils.py b/megatron/training/argument_utils.py index 70f26c64d56..e30e14cfbab 100644 --- a/megatron/training/argument_utils.py +++ b/megatron/training/argument_utils.py @@ -274,6 +274,90 @@ def _get_field_docstrings(self, src_cfg_class: type) -> dict[str, str]: return field_docstrings +def _normalize_dsv4_hybrid_csa_compress_ratios(args, kw_args, pattern): + """Normalize compact HybridModel ratios into a full per-layer config list.""" + from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols + + variant = kw_args.get( + 'experimental_attention_variant', getattr(args, 'experimental_attention_variant', None) + ) + if variant != 'dsv4_hybrid': + return + + fixed_ratio_map = {Symbols.WINDOW: 0, Symbols.CSA: 4, Symbols.HCA: 128} + ratio_symbols = set(fixed_ratio_map) | {Symbols.DS_ATTENTION} + sections = pattern.split(Symbols.MTP_SEPARATOR) + layers = ''.join(section.replace(Symbols.PIPE, '') for section in sections) + attention_symbols = [symbol for symbol in layers if symbol in ratio_symbols] + compact_len = len(attention_symbols) + full_len = len(layers) + + def pad_compact_ratios(provided): + compact = [] + full = [] + compact_iter = iter(provided) + for symbol in layers: + if symbol in ratio_symbols: + ratio = next(compact_iter) + if symbol in fixed_ratio_map: + expected = fixed_ratio_map[symbol] + assert ratio == expected, ( + f"csa_compress_ratios has ratio {ratio} for hybrid symbol " + f"'{symbol}', expected {expected}." + ) + else: + assert ratio in (0, 4, 128), ( + f"csa_compress_ratios has invalid array-driven D ratio {ratio}." + ) + compact.append(ratio) + full.append(ratio) + else: + full.append(0) + return compact, full + + # Keep the CLI/checkpoint value compact while providing TransformerConfig with a + # full list indexed by global decoder/MTP layer number. + if getattr(args, 'csa_compress_ratios', None) is None: + compact_ratios = [fixed_ratio_map.get(symbol, 0) for symbol in attention_symbols] + args.csa_compress_ratios, kw_args['csa_compress_ratios'] = pad_compact_ratios( + compact_ratios + ) + return + + provided = list(args.csa_compress_ratios) + if len(provided) == compact_len: + args.csa_compress_ratios, kw_args['csa_compress_ratios'] = pad_compact_ratios(provided) + elif len(provided) == full_len: + compact = [] + for ratio, symbol in zip(provided, layers): + if symbol in ratio_symbols: + if symbol in fixed_ratio_map: + expected = fixed_ratio_map[symbol] + assert ratio == expected, ( + f"csa_compress_ratios has ratio {ratio} for hybrid symbol " + f"'{symbol}', expected {expected}." + ) + else: + assert ratio in (0, 4, 128), ( + f"csa_compress_ratios has invalid array-driven D ratio {ratio}." + ) + compact.append(ratio) + else: + assert ratio == 0, ( + "csa_compress_ratios should not pad non-attention hybrid symbol " + f"'{symbol}' with non-zero ratio {ratio}." + ) + args.csa_compress_ratios = compact + kw_args['csa_compress_ratios'] = provided + else: + raise AssertionError( + f"csa_compress_ratios length ({len(provided)}) must equal either the " + f"number of D/W/C/H attention symbols ({compact_len}) or the legacy " + f"number of all layers in the hybrid pattern ({full_len}) for pattern " + f"'{pattern}'." + ) + + def core_transformer_config_from_args(args, config_class=None): from megatron.core.activations import squared_relu from megatron.core.fusions.fused_bias_geglu import quick_gelu @@ -306,6 +390,7 @@ def core_transformer_config_from_args(args, config_class=None): kw_args['pipeline_dtype'] = args.params_dtype kw_args['batch_p2p_comm'] = not args.overlap_p2p_comm kw_args['num_moe_experts'] = args.num_experts + kw_args['actual_vocab_size'] = args.padded_vocab_size kw_args['rotary_interleaved'] = args.rotary_interleaved kw_args['num_layers_in_first_pipeline_stage']= args.decoder_first_pipeline_num_layers kw_args['num_layers_in_last_pipeline_stage']= args.decoder_last_pipeline_num_layers @@ -345,8 +430,19 @@ def core_transformer_config_from_args(args, config_class=None): if args.hybrid_layer_pattern is not None: kw_args['is_hybrid_model'] = True from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols - if Symbols.DS_ATTENTION in args.hybrid_layer_pattern: - kw_args['experimental_attention_variant'] = 'dsa' + + pattern = args.hybrid_layer_pattern + has_dsv4_csa = any( + symbol in pattern for symbol in (Symbols.WINDOW, Symbols.CSA, Symbols.HCA) + ) + has_dsa = Symbols.DS_ATTENTION in pattern + if getattr(args, 'experimental_attention_variant', None) is None: + if has_dsv4_csa: + kw_args['experimental_attention_variant'] = 'dsv4_hybrid' + elif has_dsa: + kw_args['experimental_attention_variant'] = 'dsa' + + _normalize_dsv4_hybrid_csa_compress_ratios(args, kw_args, pattern) kw_args['inference_sampling_seed'] = args.seed diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 7f84a30bae0..7549ef8490d 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -6,15 +6,16 @@ import dataclasses import json import os -from pathlib import Path import re import types +from pathlib import Path import torch +from megatron.core.model_parallel_config import _parse_pad_packed_seq_alignment +from megatron.core.msc_utils import MultiStorageClientFeature from megatron.core.rerun_state_machine import RerunStateMachine from megatron.core.transformer import TransformerConfig -from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout from megatron.core.transformer.cuda_graph_config import ( ALLOWED_INFERENCE_SCOPES, get_deprecated_cuda_graph_modules_migration, @@ -23,23 +24,24 @@ validate_deprecated_cuda_graph_modules_migration_inputs, ) from megatron.core.transformer.enums import AttnBackend, CudaGraphModule, InferenceCudaGraphScope +from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout from megatron.core.utils import ( get_torch_version, is_flashinfer_min_version, is_te_min_version, is_torch_min_version, ) +from megatron.training.argument_utils import ( # noqa: F401 # pylint: disable=unused-import + ArgumentGroupFactory, + core_transformer_config_from_args, +) from megatron.training.global_vars import set_global_variables from megatron.training.utils import ( get_device_arch_version, - update_use_dist_ckpt, print_rank_0, + update_use_dist_ckpt, warn_rank_0, ) -from megatron.core.msc_utils import MultiStorageClientFeature - -from megatron.training.argument_utils import ArgumentGroupFactory, core_transformer_config_from_args # noqa: F401 # pylint: disable=unused-import - def add_megatron_arguments(parser: argparse.ArgumentParser): @@ -79,6 +81,7 @@ def add_megatron_arguments(parser: argparse.ArgumentParser): parser = _add_msc_args(parser) parser = _add_kitchen_quantization_arguments(parser) parser = _add_sft_args(parser) + parser = _add_varlen_dataset_args(parser) parser = _add_fault_injector_args(parser) @@ -323,6 +326,15 @@ def no_rope_freq_type(x): # it's a single int but in str return int(x) + +def compress_ratios_type(x): + """Parse per-layer compression ratios for compressed sparse attention.""" + if isinstance(x, list): + return x + assert isinstance(x, str) + return _eval_pattern(x) + + def moe_freq_type(x): """Frequency between MoE layers and Dense layers. @@ -398,8 +410,9 @@ def validate_args(args, defaults={}): 'Currently only global and local checkpoints are supported' if args.non_persistent_ckpt_type == 'local': try: - from nvidia_resiliency_ext.checkpointing.local.ckpt_managers.local_manager import \ - LocalCheckpointManager + from nvidia_resiliency_ext.checkpointing.local.ckpt_managers.local_manager import ( + LocalCheckpointManager, + ) except ModuleNotFoundError as e: raise RuntimeError('nvidia_resiliency_ext is required for local checkpointing') from e @@ -719,8 +732,10 @@ def validate_args(args, defaults={}): ) from megatron.core.models.hybrid.hybrid_layer_allocation import ( - Symbols, parse_hybrid_pattern, get_hybrid_total_layer_count, + Symbols, + get_hybrid_total_layer_count, get_hybrid_total_pipeline_segment_count, + parse_hybrid_pattern, ) sep = Symbols.MTP_SEPARATOR @@ -860,8 +875,10 @@ def validate_args(args, defaults={}): args.rank ) - # Infer use of MLA from unified pattern - if args.hybrid_layer_pattern and Symbols.DS_ATTENTION in args.hybrid_layer_pattern: + # All D/C/H/W hybrid attention symbols use MLA projections. + if args.hybrid_layer_pattern and any( + symbol in args.hybrid_layer_pattern for symbol in Symbols.MLA_ATTENTION + ): args.multi_latent_attention = True # === End of hybrid layer pattern: deprecation handling and validation === @@ -1490,6 +1507,67 @@ def validate_args(args, defaults={}): if args.ckpt_format == "fsdp_dtensor": assert args.use_megatron_fsdp, "--ckpt-format fsdp_dtensor is only tested with Megatron FSDP." + _validate_varlen_dataset_args(args) + + if args.sequence_packing_scheduler is not None: + assert not args.hybrid_context_parallel, ( + "--sequence-packing-scheduler and --hybrid-context-parallel are " + "separate scheduling paths and cannot be enabled together" + ) + assert args.calculate_per_token_loss, ( + "Sequence packing requires --calculate-per-token-loss so gradients " + "do not depend on packing boundaries" + ) + args.variable_seq_lengths = True + assert args.max_seqlen_per_dp_cp_rank is not None, ( + "--max-seqlen-per-dp-cp-rank must be set when using sequence packing" + ) + packed_capacity = args.context_parallel_size * args.max_seqlen_per_dp_cp_rank + assert packed_capacity >= args.seq_length, ( + f"Packed sequence capacity ({packed_capacity}) must be at least " + f"--seq-length ({args.seq_length})" + ) + + if getattr(args, 'pad_packed_seq_alignment', None) is not None: + args.pad_packed_seq_alignment = _parse_pad_packed_seq_alignment( + args.pad_packed_seq_alignment + ) + if args.max_seqlen_per_dp_cp_rank is None: + raise ValueError( + '--max-seqlen-per-dp-cp-rank must be set when ' + '--pad-packed-seq-alignment is enabled.' + ) + if args.pad_packed_seq_alignment != 'max': + if args.pad_packed_seq_alignment <= 0: + raise ValueError( + "--pad-packed-seq-alignment must be 'max' or a positive integer." + ) + if args.pad_packed_seq_alignment > args.max_seqlen_per_dp_cp_rank: + raise ValueError( + '--pad-packed-seq-alignment must not exceed ' + f'--max-seqlen-per-dp-cp-rank ({args.max_seqlen_per_dp_cp_rank}), ' + f'got {args.pad_packed_seq_alignment}.' + ) + + if ( + args.cuda_graph_impl == "transformer_engine" + and args.sequence_packing_scheduler is not None + ): + if getattr(args, 'pad_packed_seq_alignment', None) is None: + raise ValueError( + 'THD CUDA Graph requires --pad-packed-seq-alignment to be set.' + ) + if ( + args.pad_packed_seq_alignment != 'max' + and args.pad_packed_seq_alignment != args.max_seqlen_per_dp_cp_rank + ): + raise ValueError( + "THD CUDA Graph requires --pad-packed-seq-alignment='max' or an " + "alignment equal to --max-seqlen-per-dp-cp-rank " + f"({args.max_seqlen_per_dp_cp_rank}), got " + f"{args.pad_packed_seq_alignment}." + ) + # Data blend checks assert args.mock_data + \ bool(args.data_path) + \ @@ -2082,6 +2160,7 @@ def _add_network_size_args(parser): "no_rope_freq", "moe_layer_freq", "linear_attention_freq", + "csa_compress_ratios", "moe_router_load_balancing_type", "moe_aux_loss_coeff", "cp_comm_type", @@ -2137,8 +2216,11 @@ def _add_network_size_args(parser): "barrier_with_L1_time", # args uses same var with a different name "num_moe_experts", + "actual_vocab_size", "fp8_param", "fp4_param", + # already generated by data args + "create_attention_mask_in_dataloader", # incompatible defaults in dataclass "gradient_accumulation_fusion", "overlap_p2p_comm", @@ -2148,6 +2230,8 @@ def _add_network_size_args(parser): "bias_dropout_fusion", "apply_rope_fusion", "mamba_training_ssm_states_dtype", + "sequence_packing_scheduler", + "apply_dsa_kernel_fusion", ] transformer_factory = ArgumentGroupFactory(TransformerConfig, exclude=exclude) transformer_group = transformer_factory.build_group(parser, "transformer configuration") @@ -2579,8 +2663,7 @@ def _add_rl_args(parser): return parser def _add_training_args(parser): - from megatron.training.config import TrainingConfig - from megatron.training.config import ProfilingConfig + from megatron.training.config import ProfilingConfig, TrainingConfig prof_factory = ArgumentGroupFactory(ProfilingConfig) prof_group = prof_factory.build_group(parser, "profiling") @@ -2919,6 +3002,9 @@ def _add_distributed_args(parser): 'all layers will share the same communication type. Users can also ' 'specify separated types for each layer like ' '--cp-comm-type p2p p2p a2a a2a a2a+p2p a2a+p2p') + group.add_argument('--sequence-packing-scheduler', type=str, default=None, + choices=['dp_balanced'], + help='Pack variable-length sequences across DP x CP ranks.') group.add_argument('--fake-process-group', action='store_true', default=False, help='If set, initialize with fake distributed process group and all distributed communication operations will be skipped. \ This is quite useful for profiling memory usage of distributed training with just one GPU. \ @@ -3256,6 +3342,11 @@ def _add_mla_args(parser): help="Mscale for YaRN RoPE in multi-latent attention.") group.add_argument('--mscale-all-dim', type=float, default=0.0, help="Mscale all dimensions for YaRN RoPE in multi-latent attention.") + group.add_argument('--o-groups', type=int, default=8, + help="Number of groups for grouped low-rank output projection (wo_a).") + group.add_argument('--o-lora-rank', type=int, default=1024, + help="Low-rank dimension per group for grouped output (wo_a). " + "Used when o-groups > 0.") group.add_argument('--cache-mla-latents', action='store_true', default=False, help="If set caches the mla down projected latents with mla flash decode.") group.add_argument( @@ -3280,6 +3371,22 @@ def _add_experimental_attention_variant_args(parser): 'where 1 indicates an LA layer and 0 indicates a SDPA layer. ' 'Examples: "([0]+[1]*23)": 1 SDPA layer followed by 23 LA layers, ' '"([1]*3+[0]*2)*2": Three LA layers followed by two SDPA layers, repeated twice.') + group.add_argument( + '--csa-compress-ratios', + type=compress_ratios_type, + default=None, + help='Per-layer compress ratios for compressed sparse attention. ' + 'Accepts a Python list expression such as "[0,0,4,128,4,128]" or ' + '"([0]+[4,128]*2)*3". Valid values are 0, 4, and 128, and the ' + 'list length must equal num_layers plus mtp_num_layers.', + ) + group.add_argument( + '--no-dsa-kernel-fusion', + action='store_false', + help='Disable fused DSA sparse-attention kernels (FlashMLA + cuDNN DSA) ' + 'and fall back to unfused PyTorch implementations.', + dest='apply_dsa_kernel_fusion', + ) return parser def _add_heterogeneous_args(parser): @@ -3439,7 +3546,7 @@ def _add_kitchen_quantization_arguments(parser: argparse.ArgumentParser): If kitchen isn't available, nothing to do here, return unchanged parser """ try: - from megatron.core.extensions.kitchen import KitchenSpecProvider, HAVE_KITCHEN + from megatron.core.extensions.kitchen import HAVE_KITCHEN, KitchenSpecProvider except (ImportError, ModuleNotFoundError): HAVE_KITCHEN = False @@ -3470,6 +3577,68 @@ def _add_sft_args(parser): help='SFT prompt format.') return parser + +def _add_varlen_dataset_args(parser): + group = parser.add_argument_group(title='variable-length dataset') + group.add_argument( + '--use-varlen-dataset', + action='store_true', + help='Use variable-length raw-text pretraining data.', + ) + group.add_argument( + '--varlen-sbhd-validation', + action='store_true', + help='Use a dense-model fixed-shape reference padded to --seq-length.', + ) + group.add_argument( + '--varlen-mock-dataset-config-json', + type=str, + default=None, + help='Inline JSON or a JSON file configuring synthetic variable-length samples.', + ) + return parser + + +def _validate_varlen_dataset_args(args): + assert ( + not args.varlen_sbhd_validation or args.use_varlen_dataset + ), "--varlen-sbhd-validation requires --use-varlen-dataset" + assert ( + args.varlen_mock_dataset_config_json is None or args.use_varlen_dataset + ), "--varlen-mock-dataset-config-json requires --use-varlen-dataset" + assert ( + args.varlen_mock_dataset_config_json is None or args.mock_data + ), "--varlen-mock-dataset-config-json requires --mock-data" + if not args.use_varlen_dataset: + return + + assert not args.sft, "--use-varlen-dataset and --sft are mutually exclusive" + assert not args.fim_data, "--use-varlen-dataset and --fim-data are mutually exclusive" + assert not args.hybrid_context_parallel, ( + "--use-varlen-dataset uses the dp_balanced static-CP scheduler and cannot use " + "--hybrid-context-parallel" + ) + assert not args.reset_position_ids, "--use-varlen-dataset does not support reset position ids" + assert not args.reset_attention_mask, ( + "--use-varlen-dataset does not support reset attention masks" + ) + assert not args.dataloader_inter_document_masking, ( + "--use-varlen-dataset does not support inter-document masking" + ) + args.create_attention_mask_in_dataloader = False + if args.varlen_sbhd_validation: + assert ( + args.sequence_packing_scheduler is None + ), "--varlen-sbhd-validation cannot use a sequence packing scheduler" + assert not args.mock_data, "--varlen-sbhd-validation is only supported with real datasets" + assert args.num_experts is None, ( + "--varlen-sbhd-validation currently supports dense models only; " + "MoE requires physical padding-mask propagation" + ) + elif args.sequence_packing_scheduler is None: + args.sequence_packing_scheduler = "dp_balanced" + + def _add_logits_distillation_args(parser): group = parser.add_argument_group(title='Logits Distillation') diff --git a/megatron/training/datasets/data_samplers.py b/megatron/training/datasets/data_samplers.py index 296acc97941..80fc8327af9 100644 --- a/megatron/training/datasets/data_samplers.py +++ b/megatron/training/datasets/data_samplers.py @@ -11,7 +11,6 @@ from megatron.core import mpu from megatron.core.datasets.utils import Split - from megatron.training import get_args from megatron.training.dist_signal_handler import DistributedSignalHandler @@ -98,8 +97,11 @@ def close_nvidia_fds(): worker_init_fn if args.num_workers > 0 else None ) # Torch dataloader. - if args.hybrid_context_parallel: - extra_kwargs = {"collate_fn": lambda x: x,} + if ( + args.hybrid_context_parallel + or getattr(args, "sequence_packing_scheduler", None) is not None + ): + extra_kwargs = {"collate_fn": lambda x: x} else: extra_kwargs = {} return torch.utils.data.DataLoader( diff --git a/megatron/training/datasets/utils.py b/megatron/training/datasets/utils.py new file mode 100644 index 00000000000..4a84a82d804 --- /dev/null +++ b/megatron/training/datasets/utils.py @@ -0,0 +1,27 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Shared utilities for training-side dataset helpers.""" + +import json +import os +from typing import Any, Optional + + +def load_json_arg(spec: Optional[str]) -> Optional[Any]: + """Parse a CLI JSON argument that may be either a JSON literal or a path + to a JSON file. + + The argument is interpreted as a file path when ``spec`` points to an + existing regular file on the local filesystem; otherwise it is parsed as + a JSON literal string. Returns ``None`` when ``spec`` itself is ``None``, + so callers can use it transparently for optional CLI flags. + + Used by dataset configuration flags that accept either an inline JSON + snippet or a path to a file containing the same JSON document. + """ + if spec is None: + return None + if os.path.isfile(spec): + with open(spec, "r") as f: + return json.load(f) + return json.loads(spec) diff --git a/megatron/training/datasets/varlen_dataset.py b/megatron/training/datasets/varlen_dataset.py new file mode 100644 index 00000000000..cd5f674498d --- /dev/null +++ b/megatron/training/datasets/varlen_dataset.py @@ -0,0 +1,301 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Variable-length raw-text datasets for packed THD pretraining. + +The real-data path accepts HuggingFace datasets, local parquet files, and local +JSON/JSONL files with a ``text`` column. The mock path generates samples from +file-based or lognormal sequence-length distributions. Both paths emit one +unpacked sample at a time for the sequence-packing scheduler. +""" + +import os +from dataclasses import dataclass +from typing import Any, Dict, Optional + +import numpy as np +import torch + +from megatron.core.datasets.gpt_dataset import GPTDatasetConfig +from megatron.core.datasets.megatron_dataset import LowLevelDataset, MegatronDataset + + +@dataclass +class VarlenDatasetConfig(GPTDatasetConfig): + """Configuration for variable-length real and mock pretraining datasets.""" + + mock_dataset_config: Optional[Dict[str, Any]] = None + sbhd_validation: bool = False + + def __post_init__(self) -> None: + super().__post_init__() + assert not self.hybrid_context_parallel, ( + "VarlenDataset uses the dp_balanced static-CP scheduler and cannot use " + "hybrid context parallelism" + ) + + +def _looks_like_hf_id(path: str) -> bool: + """Return whether ``path`` looks like a HuggingFace dataset identifier.""" + if not path or os.path.exists(path) or path.startswith(("/", "./", "../")): + return False + return "/" in path + + +def _raw_text_loader(sample: Dict[str, Any]) -> str: + """Return and validate a pretraining sample's ``text`` field.""" + text = sample.get("text", "") + if text is None: + text = "" + if not isinstance(text, str): + raise ValueError( + "VarlenDataset requires the 'text' field to be a string, " + f"got {type(text).__name__}." + ) + return text + + +class VarlenLowLevelDataset: + """Load variable-length pretraining text from HF, parquet, JSON, or JSONL.""" + + def __init__(self, dataset_path: str) -> None: + try: + from datasets import Dataset, load_dataset + except ImportError as exc: + raise ImportError( + "VarlenDataset requires the `datasets` library (pip install datasets)." + ) from exc + + if _looks_like_hf_id(dataset_path): + self.dataset = load_dataset(dataset_path, split="train") + elif dataset_path.endswith(".parquet"): + self.dataset = load_dataset("parquet", data_files=dataset_path, split="all") + else: + try: + import pandas as pd + except ImportError as exc: + raise ImportError( + "VarlenDataset requires `pandas` to load local JSON/JSONL files " + "(pip install pandas)." + ) from exc + dataframe = pd.read_json( + dataset_path, lines=not dataset_path.lower().endswith(".json") + ) + self.dataset = Dataset.from_pandas(dataframe, preserve_index=False) + + if "text" not in self.dataset.column_names: + raise ValueError( + "VarlenDataset requires a raw pretraining 'text' column, " + f"got {sorted(self.dataset.column_names)}." + ) + + @property + def schema_name(self) -> str: + """Return the single supported real-data schema.""" + return "pretrain-text" + + def __len__(self) -> int: + return len(self.dataset) + + def __getitem__(self, idx: int) -> str: + return _raw_text_loader(self.dataset[idx]) + + +class VarlenDataset(MegatronDataset): + """Variable-length raw-text dataset consumed by the packing scheduler.""" + + @staticmethod + def numel_low_level_dataset(low_level_dataset: LowLevelDataset) -> int: + return len(low_level_dataset) + + @staticmethod + def build_low_level_dataset( + dataset_path: str, config: VarlenDatasetConfig + ) -> LowLevelDataset: + return VarlenLowLevelDataset(dataset_path) + + def __len__(self) -> int: + return self.num_samples + + def _calculate_padding_divisor(self) -> int: + """Return the per-sample alignment required before DP/CP packing.""" + cp_size = self.config.context_parallel_size or 1 + cp_pad = cp_size * 2 if cp_size > 1 else 1 + sp_pad = self.config.sequence_parallel_size or 1 + return cp_pad * sp_pad + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + tokenizer = self.config.tokenizer + max_len = self.config.sequence_length + eod = tokenizer.eod + pad = tokenizer.pad if tokenizer.pad is not None else eod + assert eod is not None, "VarlenDataset requires an EOD/EOS token id." + assert not self.config.reset_position_ids + assert not self.config.create_attention_mask and not self.config.reset_attention_mask + + text = self.dataset[int(self.indices[idx % len(self.indices)])] + if not isinstance(text, str): + raise ValueError( + "VarlenDataset expects raw pretraining text, " + f"got {type(text).__name__}." + ) + tokens_list = list(tokenizer.tokenize(text)) + targets_list = list(tokens_list) + + if not tokens_list: + tokens_list = [eod, eod] + targets_list = [eod, eod] + + if len(tokens_list) > max_len + 1: + tokens_list = tokens_list[: max_len + 1] + targets_list = targets_list[: max_len + 1] + if len(tokens_list) == max_len + 1 and tokens_list[-1] != eod: + tokens_list[-1] = eod + targets_list[-1] = eod + if tokens_list[-1] != eod: + tokens_list.append(eod) + targets_list.append(eod) + + valid_len = len(tokens_list) - 1 + + if self.config.sbhd_validation: + pad_len = max_len + 1 - len(tokens_list) + if pad_len > 0: + tokens_list.extend([pad] * pad_len) + targets_list.extend([pad] * pad_len) + assert len(tokens_list) == max_len + 1 + input_ids = torch.tensor(tokens_list[:-1], dtype=torch.int64) + labels = torch.tensor(targets_list[1:], dtype=torch.int64) + loss_mask = torch.ones(max_len, dtype=torch.float32) + loss_mask[valid_len:] = 0.0 + if self.config.eod_mask_loss: + loss_mask[input_ids == eod] = 0.0 + return { + "tokens": input_ids, + "labels": labels, + "loss_mask": loss_mask, + "position_ids": torch.arange(max_len, dtype=torch.int64), + } + + original_seq_len = len(tokens_list) - 1 + padding_divisor = self._calculate_padding_divisor() + remainder = original_seq_len % padding_divisor + if remainder: + pad_len = padding_divisor - remainder + tokens_list.extend([pad] * pad_len) + targets_list.extend([pad] * pad_len) + padded_seq_len = len(tokens_list) - 1 + + input_ids = torch.tensor(tokens_list[:-1], dtype=torch.int64) + labels = torch.tensor(targets_list[1:], dtype=torch.int64) + loss_mask = torch.ones(padded_seq_len, dtype=torch.float32) + loss_mask[valid_len:] = 0.0 + if self.config.eod_mask_loss: + loss_mask[input_ids == eod] = 0.0 + + return { + "tokens": input_ids, + "labels": labels, + "loss_mask": loss_mask, + "position_ids": torch.arange(padded_seq_len, dtype=torch.int64), + "original_seq_len": torch.tensor([original_seq_len], dtype=torch.int32), + "padded_seq_len": torch.tensor([padded_seq_len], dtype=torch.int32), + } + + +class MockVarlenLowLevelDataset: + """Generate mock token arrays from file-based or lognormal lengths.""" + + seed: int = 0 + size: int = 1_000_000 + + def __init__(self, mode: str, **kwargs) -> None: + np.random.seed(self.seed) + if mode == "file": + try: + import pandas as pd + except ImportError as exc: + raise ImportError( + "MockVarlenDataset file mode requires pandas (pip install pandas)." + ) from exc + self.sequence_lengths = np.asarray(pd.read_csv(kwargs["path"])).flatten() + self.size = len(self.sequence_lengths) + elif mode == "distribution": + if kwargs["type"] != "lognormal": + raise ValueError(f"Unsupported distribution type {kwargs['type']}") + sigma = kwargs["lognormal_sigma"] + mean = kwargs["mean_seq_len"] + mu = np.log(mean) - sigma**2 / 2 + samples = np.random.lognormal(mu, sigma, self.size) + self.sequence_lengths = np.clip( + samples, kwargs["min_seq_len"], kwargs["max_seq_len"] + ).astype(int) + else: + raise ValueError(f"Unsupported mode '{mode}', must be 'file' or 'distribution'") + + def __len__(self) -> int: + return self.size + + def __getitem__(self, idx: int) -> np.ndarray: + length = self.sequence_lengths[idx % self.size] + return np.arange(1, length, dtype=np.int64) + + +class MockVarlenDataset(VarlenDataset): + """Mock variable-length dataset for the packed THD benchmark path.""" + + @staticmethod + def build_low_level_dataset( + dataset_path: str, config: VarlenDatasetConfig + ) -> LowLevelDataset: + mock_config = config.mock_dataset_config + if mock_config is None: + mock_config = { + "mode": "distribution", + "type": "lognormal", + "min_seq_len": config.sequence_length // 2, + "max_seq_len": config.sequence_length, + "mean_seq_len": config.sequence_length // 4 * 3, + "lognormal_sigma": 1.1, + } + return MockVarlenLowLevelDataset(**mock_config) + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + tokenizer = self.config.tokenizer + max_len = self.config.sequence_length + eod = tokenizer.eod + pad = tokenizer.pad if tokenizer.pad is not None else eod + + tokens_list = self.dataset[int(self.indices[idx % len(self.indices)])].tolist() + if not tokens_list: + tokens_list.append(eod) + tokens_list.append(eod) + targets_list = list(tokens_list) + + if len(tokens_list) > max_len + 1: + tokens_list = tokens_list[:max_len] + [eod] + targets_list = targets_list[:max_len] + [eod] + original_seq_len = len(tokens_list) - 1 + + padding_divisor = self._calculate_padding_divisor() + remainder = original_seq_len % padding_divisor + if remainder: + pad_len = padding_divisor - remainder + tokens_list.extend([pad] * pad_len) + targets_list.extend([pad] * pad_len) + padded_seq_len = len(tokens_list) - 1 + + input_ids = torch.tensor(tokens_list[:-1], dtype=torch.int64) + labels = torch.tensor(targets_list[1:], dtype=torch.int64) + loss_mask = torch.ones(padded_seq_len, dtype=torch.float32) + loss_mask[original_seq_len:] = 0.0 + if self.config.eod_mask_loss: + loss_mask[input_ids == eod] = 0.0 + + return { + "tokens": input_ids, + "labels": labels, + "loss_mask": loss_mask, + "position_ids": torch.arange(padded_seq_len, dtype=torch.int64), + "original_seq_len": torch.tensor([original_seq_len], dtype=torch.int32), + "padded_seq_len": torch.tensor([padded_seq_len], dtype=torch.int32), + } diff --git a/megatron/training/models/hybrid.py b/megatron/training/models/hybrid.py index 287ca8ec2a3..8192dcae615 100644 --- a/megatron/training/models/hybrid.py +++ b/megatron/training/models/hybrid.py @@ -63,7 +63,7 @@ class HybridModelConfig(ModelConfig): rotary_base: int = 10000 seq_len_interpolation_factor: float | None = None make_vocab_size_divisible_by: int = 128 - hybrid_stack_spec: ModuleSpec | None = None + hybrid_stack_spec: ModuleSpec | Callable[[TransformerConfig], ModuleSpec] | None = None vocab_size: int | None = None should_pad_vocab: bool = False @@ -159,6 +159,8 @@ def build_model( ) else: hybrid_stack_spec = default_hybrid_stack_spec + elif not isinstance(hybrid_stack_spec, ModuleSpec) and callable(hybrid_stack_spec): + hybrid_stack_spec = hybrid_stack_spec(self._model_config.transformer) assert self._model_config.vocab_size is not None, "vocab_size must be configured before calling build_model()" if self._model_config.should_pad_vocab: diff --git a/megatron/training/training.py b/megatron/training/training.py index 0e4c1ea6a23..7ab7a6e282d 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -42,7 +42,7 @@ # First-party. from megatron.core import mpu, nccl_allocator, tensor_parallel -from megatron.core.datasets.data_schedule import HybridCPDataLoaderWrapper +from megatron.core.datasets.data_schedule import HybridCPDataLoaderWrapper, wrap_data_iterator from megatron.core.distributed import DistributedDataParallel as DDP from megatron.core.distributed import ( DistributedDataParallelConfig, @@ -254,6 +254,7 @@ # never call ``update_*`` so the flag stays ``False`` and no collective fires. _seqlen_stats_in_iteration: Optional[torch.Tensor] = None _seqlen_stats_active: bool = False +_seqlen_stats_are_global: bool = False # Only report memory for first 3 checkpoint saves. num_checkpoints_memory_reported = 0 @@ -315,7 +316,7 @@ def update_seqlen_stats_from_cu_seqlens(cu_seqlens): the all-reduce; BSHD callers that never invoke this function leave the flag at ``False`` and pay zero collective cost. """ - global _seqlen_stats_in_iteration, _seqlen_stats_active + global _seqlen_stats_in_iteration, _seqlen_stats_active, _seqlen_stats_are_global if cu_seqlens is None or cu_seqlens.numel() < 2: return # Pin the accumulator to the current CUDA device when available so the @@ -334,6 +335,25 @@ def update_seqlen_stats_from_cu_seqlens(cu_seqlens): _seqlen_stats_in_iteration[0] += seqlens.sum() _seqlen_stats_in_iteration[1] += (seqlens * seqlens).sum() _seqlen_stats_active = True + _seqlen_stats_are_global = False + + +def set_seqlen_stats_in_iteration(total_real_tokens, seqlen_squared_sum): + """Seed per-iteration THD FLOPs stats that were already computed globally.""" + global _seqlen_stats_in_iteration, _seqlen_stats_active, _seqlen_stats_are_global + if total_real_tokens is None or seqlen_squared_sum is None: + return + if _seqlen_stats_in_iteration is None: + device = ( + torch.device(f'cuda:{torch.cuda.current_device()}') + if torch.cuda.is_available() + else 'cpu' + ) + _seqlen_stats_in_iteration = torch.zeros(2, dtype=torch.float64, device=device) + _seqlen_stats_in_iteration[0] = float(total_real_tokens) + _seqlen_stats_in_iteration[1] = float(seqlen_squared_sum) + _seqlen_stats_active = True + _seqlen_stats_are_global = True def consume_seqlen_stats_in_iteration() -> Tuple[Optional[float], Optional[float]]: @@ -354,20 +374,23 @@ def consume_seqlen_stats_in_iteration() -> Tuple[Optional[float], Optional[float fuse onto one device tensor and consume issues a single 2-element all-reduce. - Sync cost: exactly ONE all-reduce of a 2-element ``float64`` tensor and ONE - host sync (``tolist()``). Skipped entirely when the flag is ``False``. + Sync cost: at most one all-reduce of a 2-element ``float64`` tensor and one + host sync (``tolist()``). Scheduler-provided global values skip the all-reduce; + inactive BSHD iterations skip both. - All ranks within one DP group accumulated identical values (``cu_seqlens`` is - replicated across TP/CP/PP); the world all-reduce therefore overcounts by a - factor of ``TP * CP * PP``, which we divide out. + For locally accumulated values, all ranks within one DP group see identical + ``cu_seqlens``. The world all-reduce therefore overcounts by ``TP * CP * PP``, + which we divide out. """ - global _seqlen_stats_in_iteration, _seqlen_stats_active + global _seqlen_stats_in_iteration, _seqlen_stats_active, _seqlen_stats_are_global if not _seqlen_stats_active: # BSHD path: never allocated the tensor; tell the caller to use the # closed-form defaults. return None, None t = _seqlen_stats_in_iteration - if torch.distributed.is_initialized() and mpu.model_parallel_is_initialized(): + if _seqlen_stats_are_global: + dedup = 1 + elif torch.distributed.is_initialized() and mpu.model_parallel_is_initialized(): torch.distributed.all_reduce(t) tp_size = max(mpu.get_tensor_model_parallel_world_size(), 1) cp_size = max(mpu.get_context_parallel_world_size(), 1) @@ -383,9 +406,108 @@ def consume_seqlen_stats_in_iteration() -> Tuple[Optional[float], Optional[float # iterations reuse it without reallocating. t.zero_() _seqlen_stats_active = False + _seqlen_stats_are_global = False return total_real_tokens / dedup, seqlen_squared_sum / dedup +def _dsv4_hybrid_self_attention_flops( + *, + hidden_size, + num_attention_heads, + v_head_dim, + q_lora_rank, + o_groups, + o_lora_rank, + csa_window_size, + seq_length, + n_layers_r0, + n_layers_r4, + n_layers_r128, + dsa_indexer_n_heads, + dsa_indexer_head_dim, + dsa_indexer_topk, +): + """Return token-linear and L-squared DSv4 HybridModel attention coefficients. + + Ratio-zero layers use window attention, ratio-four layers add learned sparse + selection over compressed KV, and ratio-128 layers attend to all compressed + KV. Expansion factors for fused multiply-add and forward/backward are left + to the caller. + """ + n_attn_layers = n_layers_r0 + n_layers_r4 + n_layers_r128 + + # DSv4 uses a joint hidden-to-v_head_dim KV projection and grouped low-rank + # output projection, rather than generic MLA's dense output projection. + q_term = q_lora_rank * (hidden_size + num_attention_heads * v_head_dim + 1) + kv_term = hidden_size * v_head_dim + v_head_dim + o_term = ( + num_attention_heads * v_head_dim * o_lora_rank + + o_groups * o_lora_rank * hidden_size + ) + mla_projection_term = (q_term + kv_term + o_term) * n_attn_layers + + sparse_r0 = ( + n_layers_r0 * num_attention_heads * csa_window_size * v_head_dim * 2 + ) + sparse_r128_window = ( + n_layers_r128 * num_attention_heads * csa_window_size * v_head_dim * 2 + ) + sparse_r128_core = n_layers_r128 * num_attention_heads * v_head_dim / 128 + + # The main compressor has wkv and wgate projections. Ratio-four layers use + # two overlapping compressor windows; ratio-128 layers use one. + compressor_term = ( + n_layers_r4 * hidden_size * (2 * v_head_dim) * 2 + + n_layers_r128 * hidden_size * v_head_dim * 2 + ) + + if n_layers_r4: + assert dsa_indexer_n_heads is not None + assert dsa_indexer_head_dim is not None + assert dsa_indexer_topk is not None + + # Packed iterations can contain shorter subsequences, but the scheduler + # exposes only aggregate sum(L) and sum(L^2). Use configured seq_length + # for the ratio-four effective-top-k approximation. + effective_topk = min(dsa_indexer_topk, seq_length // 4) + average_compressed_tokens = effective_topk * ( + 1 - effective_topk * 4 / (2 * seq_length) + ) + sparse_r4 = ( + n_layers_r4 + * num_attention_heads + * (csa_window_size + average_compressed_tokens) + * v_head_dim + * 2 + ) + indexer_token_term = ( + n_layers_r4 * hidden_size * (2 * dsa_indexer_head_dim) * 2 + + n_layers_r4 + * q_lora_rank + * dsa_indexer_n_heads + * dsa_indexer_head_dim + + n_layers_r4 * hidden_size * dsa_indexer_n_heads + ) + indexer_core_term = ( + n_layers_r4 * dsa_indexer_n_heads * dsa_indexer_head_dim / 4 + ) + else: + sparse_r4 = 0 + indexer_token_term = 0 + indexer_core_term = 0 + + token_linear = ( + mla_projection_term + + sparse_r0 + + sparse_r4 + + sparse_r128_window + + compressor_term + + indexer_token_term + ) + core = sparse_r128_core + indexer_core_term + return token_linear, core + + def num_floating_point_operations( args, batch_size, @@ -523,12 +645,49 @@ def hybrid_flops(total_tokens, seqlen_squared_sum, hidden_size, gdn_qk_head_dim=128, gdn_v_head_dim=128, gdn_num_qk_heads=16, gdn_num_v_heads=32, gdn_conv_kernel_dim=4, - vocab_size=256000, mtp_num_layers=0): + vocab_size=256000, mtp_num_layers=0, + experimental_attention_variant=None, + q_lora_rank=None, v_head_dim=None, + o_groups=None, o_lora_rank=None, + csa_window_size=None, seq_length=None, + dsv4_n_layers_r0=0, dsv4_n_layers_r4=0, dsv4_n_layers_r128=0, + dsa_indexer_n_heads=None, dsa_indexer_head_dim=None, + dsa_indexer_topk=None): """Calculate total FLOPs for the hybrid model.""" + if experimental_attention_variant == "dsv4_hybrid": + token_term, core_term = _dsv4_hybrid_self_attention_flops( + hidden_size=hidden_size, + num_attention_heads=num_attn_heads, + v_head_dim=v_head_dim, + q_lora_rank=q_lora_rank, + o_groups=o_groups, + o_lora_rank=o_lora_rank, + csa_window_size=csa_window_size, + seq_length=seq_length, + n_layers_r0=dsv4_n_layers_r0, + n_layers_r4=dsv4_n_layers_r4, + n_layers_r128=dsv4_n_layers_r128, + dsa_indexer_n_heads=dsa_indexer_n_heads, + dsa_indexer_head_dim=dsa_indexer_head_dim, + dsa_indexer_topk=dsa_indexer_topk, + ) + # FMA expansion is applied here; the final *3 accounts for + # forward, weight-gradient, and data-gradient work. + attention_flops = 2 * ( + token_term * total_tokens + core_term * seqlen_squared_sum + ) + else: + attention_flops = num_attn_layers * attn_layer_flops( + total_tokens, + seqlen_squared_sum, + hidden_size, + num_attn_heads, + gqa, + gqa_groups, + kv_channels, + ) flops_fwd = ( - num_attn_layers * attn_layer_flops(total_tokens, seqlen_squared_sum, - hidden_size, num_attn_heads, gqa, - gqa_groups, kv_channels) + + attention_flops + num_mlp_layers * mlp_layer_flops(total_tokens, hidden_size, mlp_expansion, swiglu) + num_mamba_layers * mamba_layer_flops(total_tokens, hidden_size, @@ -541,6 +700,8 @@ def hybrid_flops(total_tokens, seqlen_squared_sum, hidden_size, gdn_qk_head_dim, gdn_v_head_dim, gdn_num_qk_heads, gdn_num_v_heads, gdn_conv_kernel_dim) + + # MTP eh_norm, final_norm, and eh projection. + (2 * mtp_num_layers * (3 * hidden_size + 2 * hidden_size**2) * total_tokens) + (2 * total_tokens * hidden_size * vocab_size * (1 + mtp_num_layers)) # logits computation ) return flops_fwd * 3 @@ -853,11 +1014,26 @@ def transformer_flops(): Symbols, get_hybrid_layer_counts, ) + layer_counts = get_hybrid_layer_counts(args.hybrid_layer_pattern) num_mamba_layers, num_gdn_layers, num_attn_layers, num_mlp_layers, num_moe_layers = ( itemgetter(Symbols.MAMBA, Symbols.GDN, Symbols.ATTENTION, Symbols.MLP, Symbols.MOE)( - get_hybrid_layer_counts(args.hybrid_layer_pattern) + layer_counts ) ) + dsv4_n_layers_r0 = 0 + dsv4_n_layers_r4 = 0 + dsv4_n_layers_r128 = 0 + if args.experimental_attention_variant == "dsv4_hybrid": + dsv4_n_layers_r0 = layer_counts[Symbols.WINDOW] + dsv4_n_layers_r4 = layer_counts[Symbols.CSA] + dsv4_n_layers_r128 = layer_counts[Symbols.HCA] + dsv4_num_attention_layers = ( + dsv4_n_layers_r0 + dsv4_n_layers_r4 + dsv4_n_layers_r128 + ) + assert ( + num_attn_layers == 0 and layer_counts[Symbols.DS_ATTENTION] == 0 + ), "dsv4_hybrid supports only Window/CSA/HCA attention layers." + num_attn_layers = dsv4_num_attention_layers mtp_num_layers = args.mtp_num_layers if mtp_num_layers is None: @@ -895,6 +1071,19 @@ def transformer_flops(): gdn_conv_kernel_dim=args.linear_conv_kernel_dim or 4, vocab_size=args.padded_vocab_size, mtp_num_layers=mtp_num_layers, + experimental_attention_variant=args.experimental_attention_variant, + q_lora_rank=args.q_lora_rank, + v_head_dim=args.v_head_dim, + o_groups=getattr(args, "o_groups", None), + o_lora_rank=getattr(args, "o_lora_rank", None), + csa_window_size=getattr(args, "csa_window_size", None), + seq_length=args.seq_length, + dsv4_n_layers_r0=dsv4_n_layers_r0, + dsv4_n_layers_r4=dsv4_n_layers_r4, + dsv4_n_layers_r128=dsv4_n_layers_r128, + dsa_indexer_n_heads=getattr(args, "dsa_indexer_n_heads", None), + dsa_indexer_head_dim=getattr(args, "dsa_indexer_head_dim", None), + dsa_indexer_topk=getattr(args, "dsa_indexer_topk", None), ) else: # Compute standard Transformer model FLOPs. @@ -2297,6 +2486,7 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch """ args = get_args() timers = get_timers() + scheduled_num_microbatches = get_num_microbatches() rerun_state_machine = get_rerun_state_machine() save_params_in_this_iteration = (args.save_params_interval is not None and @@ -2309,7 +2499,8 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch (iteration + 1) % args.save_wgrads_interval == 0) save_dgrads_in_this_iteration = (args.save_dgrads_interval is not None and (iteration + 1) % args.save_dgrads_interval == 0) - while rerun_state_machine.should_run_forward_backward(data_iterator): + source_data_iterator = data_iterator + while rerun_state_machine.should_run_forward_backward(source_data_iterator): # Set grad to zero. for model_chunk in model: model_chunk.zero_grad_buffer() @@ -2351,6 +2542,26 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch if isinstance(optim_instance, DistributedOptimizer): optim_instance._copy_main_params_to_param_buffer() + if getattr(config, "sequence_packing_scheduler", None) is not None: + scheduler_pg_collection = get_attr_wrapped_model(model[0], "pg_collection") + assert isinstance(scheduler_pg_collection, ProcessGroupCollection), ( + "sequence packing requires the model to expose a ProcessGroupCollection" + ) + ( + scheduled_data_iterator, + scheduled_num_microbatches, + total_real_tokens_in_batch, + seqlen_squared_sum_in_batch, + ) = wrap_data_iterator( + source_data_iterator, + config, + get_num_microbatches(), + pg_collection=scheduler_pg_collection, + ) + else: + scheduled_data_iterator = source_data_iterator + scheduled_num_microbatches = get_num_microbatches() + # Forward pass. if save_activations_in_this_iteration: enable_activation_logging(model, args.save) @@ -2360,9 +2571,9 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch enable_dgrad_logging(model, args.save) losses_reduced = forward_backward_func( forward_step_func=forward_step_func, - data_iterator=data_iterator, + data_iterator=scheduled_data_iterator, model=model, - num_microbatches=get_num_microbatches(), + num_microbatches=scheduled_num_microbatches, seq_length=args.seq_length, micro_batch_size=args.micro_batch_size, decoder_seq_length=args.decoder_seq_length, @@ -2372,6 +2583,15 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch p2p_communicator=p2p_communicator, pg_collection=pg_collection, ) + if getattr(config, "sequence_packing_scheduler", None) is not None: + # The HybridModel forward path also records local cu_seqlens stats. + # Overwrite those after forward/backward with the scheduler's + # already-global values so consume() neither double-counts nor + # issues another all-reduce. + set_seqlen_stats_in_iteration( + total_real_tokens_in_batch, + seqlen_squared_sum_in_batch, + ) if save_activations_in_this_iteration: save_activations(iteration + 1) disable_activation_logging() @@ -2411,7 +2631,17 @@ def _save_state_dict(attr_name, label): should_checkpoint, should_exit, exit_code = rerun_state_machine.should_checkpoint_and_exit() if should_exit: - return {}, True, should_checkpoint, should_exit, exit_code, None, None, 0 + return ( + {}, + True, + should_checkpoint, + should_exit, + exit_code, + None, + None, + 0, + scheduled_num_microbatches, + ) # Empty unused memory. if args.empty_unused_memory_level >= 1: @@ -2505,8 +2735,19 @@ def _save_state_dict(attr_name, label): grad_norm, num_zeros_in_grad, log_max_attention_logit, + scheduled_num_microbatches, ) - return {}, skipped_iter, should_checkpoint, should_exit, exit_code, grad_norm, num_zeros_in_grad, log_max_attention_logit + return ( + {}, + skipped_iter, + should_checkpoint, + should_exit, + exit_code, + grad_norm, + num_zeros_in_grad, + log_max_attention_logit, + scheduled_num_microbatches, + ) def training_log( @@ -2525,6 +2766,7 @@ def training_log( is_first_iteration=False, seqlen_squared_sum_in_batch: float | None = None, total_real_tokens_in_batch: float | None = None, + num_microbatches: int | None = None, ): """Log training information such as losses, timing, ....""" args = get_args() @@ -2694,7 +2936,7 @@ def training_log( # Log MoE metrics. moe_log_string = "" if args.num_experts is not None: - moe_loss_scale = 1 / get_num_microbatches() + moe_loss_scale = 1 / (num_microbatches or get_num_microbatches()) track_names = [] if "aux_loss" in args.moe_router_load_balancing_type: track_names.append("load_balancing_loss") @@ -2705,14 +2947,24 @@ def training_log( if args.moe_z_loss_coeff is not None: track_names.append("z_loss") + moe_layer_freq = args.moe_layer_freq + mtp_num_layers = args.mtp_num_layers if is_hybrid_model(args): - from operator import itemgetter - - from megatron.core.ssm.mamba_hybrid_layer_allocation import ( + from megatron.core.models.hybrid.hybrid_layer_allocation import ( Symbols, - get_hybrid_layer_counts, + parse_hybrid_pattern, ) - layers = itemgetter(Symbols.MOE)(get_hybrid_layer_counts(args.hybrid_layer_pattern)) + + parsed_hybrid_pattern = parse_hybrid_pattern(args.hybrid_layer_pattern) + main_pattern = (parsed_hybrid_pattern.main_pattern or "").replace(Symbols.PIPE, "") + layers = len(main_pattern) + parsed_hybrid_pattern.mtp_num_depths + moe_layer_freq = [int(layer_type == Symbols.MOE) for layer_type in main_pattern] + moe_layer_freq.extend( + parsed_hybrid_pattern.mtp_pattern.count(Symbols.MOE) + for _ in range(parsed_hybrid_pattern.mtp_num_depths) + ) + # The full Hybrid pattern above already accounts for each MTP depth. + mtp_num_layers = None else: layers = args.num_layers @@ -2725,28 +2977,35 @@ def training_log( force_initialize=True, track_names=track_names, num_layers=layers, - moe_layer_freq=args.moe_layer_freq, - mtp_num_layers=args.mtp_num_layers, + moe_layer_freq=moe_layer_freq, + mtp_num_layers=mtp_num_layers, pg_collection=pg_collection, total_loss_dict=total_loss_dict, ) # Log MTP metrics. if args.mtp_num_layers is not None: - mtp_loss_scale = 1 / get_num_microbatches() + mtp_loss_scale = 1 / (num_microbatches or get_num_microbatches()) MTPLossLoggingHelper.track_mtp_metrics( mtp_loss_scale, iteration, writer, wandb_writer, total_loss_dict ) # Track sparse attention indexer loss. if args.dsa_indexer_loss_coeff is not None and args.dsa_indexer_loss_coeff > 0: - indexer_loss_scale = 1 / get_num_microbatches() + indexer_loss_scale = 1 / (num_microbatches or get_num_microbatches()) + assert isinstance( + pg_collection, ProcessGroupCollection + ), "DSA indexer logging requires a ProcessGroupCollection" DSAIndexerLossLoggingHelper.track_indexer_metrics( loss_scale=indexer_loss_scale, iteration=iteration, writer=writer, + pg_collection=pg_collection, wandb_writer=wandb_writer, 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 == "transformer_engine", ) # Dump memory snapshot and print metrics to stdout. @@ -3625,12 +3884,22 @@ def trace_handler(p): # Initialize CUDA Graphs helper. if args.cuda_graph_impl == "transformer_engine": + cuda_graph_sample_packed_seq_params = None + if has_rl_utils and args.perform_rl_step: + cuda_graph_sample_packed_seq_params = rl_utils.get_rl_packed_seq_params_for_cuda_graph( + seq_length=args.seq_length, + device=torch.device("cuda", torch.cuda.current_device()), + sequence_packing=args.rl_use_sequence_packing, + max_sequences_per_bin=args.rl_sequence_packing_max_sequences_per_bin, + ) cuda_graph_helper = TECudaGraphHelper( model=model, config=config, seq_length=args.seq_length, micro_batch_size=args.micro_batch_size, optimizers=[optimizer], + sample_packed_seq_params=cuda_graph_sample_packed_seq_params, + thd_sequence_length_upper_bound=_get_thd_sequence_length_upper_bound(args), ) # Run training iterations till done. @@ -3671,7 +3940,7 @@ def trace_handler(p): # Skip automatic checkpoint on microbatch changes when sequence packing is active # as it intentionally reconfigures microbatches if get_num_microbatches() != num_microbatches and iteration != 0: - if args.rl_use_sequence_packing: + if args.rl_use_sequence_packing or args.sequence_packing_scheduler is not None: print_rank_0( f"[Sequence Packing] Skipping automatic checkpoint at iteration {iteration} " f"(microbatch change: {num_microbatches} -> {get_num_microbatches()})" @@ -3709,6 +3978,9 @@ def trace_handler(p): # Completely skip iteration if needed. if (iteration + 1) in args.iterations_to_skip: + assert ( + getattr(config, "sequence_packing_scheduler", None) is None + ), "Sequence packing scheduler is not supported in skip iteration mode" # Dummy train_step to fast forward train_data_iterator. dummy_train_step(train_data_iterator) if iteration == start_iteration: @@ -3755,6 +4027,7 @@ def trace_handler(p): grad_norm = 0.0 num_zeros_in_grad = 0 max_attention_logit = None + num_microbatches = get_num_microbatches() else: ft_integration.on_training_step_start() ( @@ -3766,6 +4039,7 @@ def trace_handler(p): grad_norm, num_zeros_in_grad, max_attention_logit, + num_microbatches, ) = train_step( forward_step_func, train_data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func, iteration=iteration, pg_collection=pg_collection, @@ -3911,6 +4185,7 @@ def trace_handler(p): is_first_iteration=is_first_iteration, seqlen_squared_sum_in_batch=seqlen_squared_sum_in_batch, total_real_tokens_in_batch=total_real_tokens_in_batch, + num_microbatches=num_microbatches, ) is_first_iteration = False @@ -4151,11 +4426,31 @@ def evaluate( # Don't care about timing during evaluation config.timers = None ft_integration.on_eval_step_start() + if getattr(config, "sequence_packing_scheduler", None) is not None: + assert isinstance( + eval_pgc, ProcessGroupCollection + ), "sequence packing requires the model to expose a ProcessGroupCollection" + try: + (packed_data_iterator, scheduled_eval_num_microbatches, _, _) = ( + wrap_data_iterator( + data_iterator, + config, + eval_num_microbatches, + pg_collection=eval_pgc, + ) + ) + except StopIteration: + ft_integration.on_eval_step_end() + config.timers = get_timers() + break + else: + packed_data_iterator = data_iterator + scheduled_eval_num_microbatches = eval_num_microbatches loss_dicts = forward_backward_func( forward_step_func=forward_step_func, - data_iterator=data_iterator, + data_iterator=packed_data_iterator, model=model, - num_microbatches=eval_num_microbatches, + num_microbatches=scheduled_eval_num_microbatches, seq_length=args.seq_length, micro_batch_size=eval_micro_batch_size, decoder_seq_length=args.decoder_seq_length, @@ -4608,3 +4903,33 @@ def should_disable_forward_pre_hook(args): ) and args.overlap_param_gather ) + + +def _get_thd_sequence_length_upper_bound(args): + """Return the padded per-sample THD length upper bound for graph sizing.""" + max_sequence_length = getattr(args, "seq_length", None) + if getattr(args, "use_varlen_dataset", False): + mock_config_spec = getattr(args, "varlen_mock_dataset_config_json", None) + if mock_config_spec is not None: + from megatron.training.datasets.utils import load_json_arg + + mock_config = load_json_arg(mock_config_spec) + if isinstance(mock_config, dict) and mock_config.get("max_seq_len") is not None: + max_sequence_length = int(mock_config["max_seq_len"]) + + if max_sequence_length is None: + return None + if getattr(args, "seq_length", None) is not None: + max_sequence_length = min(int(max_sequence_length), int(args.seq_length)) + + cp_size = int(getattr(args, "context_parallel_size", 1) or 1) + cp_padding = cp_size * 2 if cp_size > 1 else 1 + sp_padding = ( + int(getattr(args, "tensor_model_parallel_size", 1) or 1) + if getattr(args, "sequence_parallel", False) + else 1 + ) + padding_granularity = cp_padding * sp_padding + return int( + math.ceil(max_sequence_length / padding_granularity) * padding_granularity + ) diff --git a/pretrain_hybrid.py b/pretrain_hybrid.py index 39bc7f30b57..06d8bb11ace 100644 --- a/pretrain_hybrid.py +++ b/pretrain_hybrid.py @@ -25,10 +25,11 @@ from hybrid_builders import hybrid_builder from megatron.core import mpu from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder +from megatron.core.datasets.data_schedule import get_batch_on_this_rank_for_sequence_packing from megatron.core.datasets.gpt_dataset import GPTDataset, GPTDatasetConfig, MockGPTDataset from megatron.core.enums import ModelType -from megatron.core.package_info import __version__ as mcore_version from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.package_info import __version__ as mcore_version from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.parallel_state import ( get_context_parallel_group, @@ -62,6 +63,12 @@ ) from megatron.training.arguments import core_transformer_config_from_args, parse_and_validate_args from megatron.training.datasets.sft_dataset import SFTDataset +from megatron.training.datasets.utils import load_json_arg +from megatron.training.datasets.varlen_dataset import ( + MockVarlenDataset, + VarlenDataset, + VarlenDatasetConfig, +) from megatron.training.training import update_seqlen_stats_from_cu_seqlens from megatron.training.utils import get_blend_and_blend_per_split, is_first_or_last_pipeline_stage from model_provider import model_provider @@ -94,7 +101,7 @@ ] -def get_batch(data_iterator, vp_stage=None): +def get_batch(data_iterator, vp_stage=None, pg_collection=None): """Generate a batch.""" args = get_args() @@ -113,12 +120,44 @@ def get_batch(data_iterator, vp_stage=None): ) is_hybrid_cp = args.hybrid_context_parallel + if args.sequence_packing_scheduler is not None: + ( + tokens, + labels, + loss_mask, + attention_mask, + position_ids, + packed_seq_params, + padding_mask, + ) = get_batch_on_this_rank_for_sequence_packing( + data_iterator, + vpp_size=config.virtual_pipeline_model_parallel_size, + mtp_on_this_rank=mtp_on_this_rank, + vp_stage=vp_stage, + pg_collection=pg_collection, + config=config, + ) + return ( + attention_mask, + None, + None, + None, + labels, + None, + loss_mask, + None, + position_ids, + tokens, + padding_mask, + packed_seq_params, + ) + if ( not is_first_or_last_pipeline_stage(vp_stage) and not mtp_on_this_rank and not has_cu_seqlens ): - return [None for _ in BATCH_KEYS] + return [None for _ in BATCH_KEYS] + [None, None] batch = {} if tp_rank == 0: @@ -162,6 +201,8 @@ def get_batch(data_iterator, vp_stage=None): batch['max_seqlen'], None, None, + None, + None, ) batch = get_batch_on_this_cp_rank( @@ -178,7 +219,7 @@ def get_batch(data_iterator, vp_stage=None): # BATCH_KEYS entry on tp_rank 0; other tp_ranks receive a fresh dict from # get_batch_on_this_tp_rank. BATCH_KEYS is already alphabetical, matching # the historical sorted(batch.keys()) order. - return [batch[key] for key in BATCH_KEYS] + return [batch[key] for key in BATCH_KEYS] + [None, None] # define spiky loss as a loss that's 10x the max loss observed @@ -284,6 +325,7 @@ def forward_step(data_iterator, model: HybridModel): data_iterator : Input data iterator model (HybridModel): The Hybrid Model """ + args = get_args() timers = get_timers() # Get the batch. @@ -291,6 +333,7 @@ def forward_step(data_iterator, model: HybridModel): with stimer(bdata=True): vp_stage = get_attr_wrapped_model(model, "vp_stage") + pg_collection = get_attr_wrapped_model(model, "pg_collection") ( attention_mask, cu_seqlens, @@ -302,12 +345,16 @@ def forward_step(data_iterator, model: HybridModel): max_seqlen, position_ids, tokens, - ) = get_batch(data_iterator, vp_stage) - - packed_seq_params = None - if cu_seqlens is not None: - # Squeeze the batch dim: the batch dict keeps cu_seqlens as (1, N) - # for consistency, but PackedSeqParams and TE expect 1-D. + padding_mask, + packed_seq_params, + ) = get_batch(data_iterator, vp_stage, pg_collection) + + if packed_seq_params is not None: + if packed_seq_params.cu_seqlens_q is not None: + update_seqlen_stats_from_cu_seqlens(packed_seq_params.cu_seqlens_q) + elif cu_seqlens is not None: + # cu_seqlens / cu_seqlens_padded carry the dataloader's batch dim (1, n). + # PackedSeqParams and TE attention expect 1-D tensors. cu_seqlens = cu_seqlens.squeeze(0) if cu_seqlens_padded is not None: cu_seqlens_padded = cu_seqlens_padded.squeeze(0) @@ -339,6 +386,7 @@ def forward_step(data_iterator, model: HybridModel): labels=labels, packed_seq_params=packed_seq_params, loss_mask=loss_mask, + padding_mask=padding_mask, ) # [ModelOpt]: model is needed to access ModelOpt distillation losses @@ -375,7 +423,7 @@ def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig: with open(args.per_dataset_sequences_path, "r") as f: sequences_per_dataset = json.load(f) - return GPTDatasetConfig( + data_args = dict( random_seed=args.seed, sequence_length=args.seq_length, blend=blend, @@ -404,6 +452,17 @@ def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig: inter_document_masking=args.dataloader_inter_document_masking, ) + if args.use_varlen_dataset: + data_args["mock_dataset_config"] = ( + load_json_arg(args.varlen_mock_dataset_config_json) + if args.varlen_mock_dataset_config_json is not None + else None + ) + data_args["sbhd_validation"] = args.varlen_sbhd_validation + return VarlenDatasetConfig(**data_args) + + return GPTDatasetConfig(**data_args) + def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None): """Build the train test and validation datasets. @@ -412,19 +471,20 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None train_val_test_num_samples : A list containing the number of samples in train test and validation. """ args = get_args() - config = core_gpt_dataset_config_from_args(args) is_packed_sequence = False if args.sft: dataset_type = SFTDataset is_packed_sequence = True # SFT always uses packed sequence + elif args.use_varlen_dataset: + dataset_type = MockVarlenDataset if args.mock_data else VarlenDataset + is_packed_sequence = not args.varlen_sbhd_validation else: - if args.mock_data: - dataset_type = MockGPTDataset - else: - dataset_type = GPTDataset + dataset_type = MockGPTDataset if args.mock_data else GPTDataset + + config = core_gpt_dataset_config_from_args(args) - print_rank_0("> building train, validation, and test datasets for GPT ...") + print_rank_0("> building train, validation, and test datasets for HybridModel ...") train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder( dataset_type, @@ -433,7 +493,7 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None config, ).build() - print_rank_0("> finished creating GPT datasets ...") + print_rank_0("> finished creating HybridModel datasets ...") return train_ds, valid_ds, test_ds diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp1_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp1_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml new file mode 100644 index 00000000000..17ed2c24201 --- /dev/null +++ b/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp1_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml @@ -0,0 +1,106 @@ +ENV_VARS: + CUDA_DEVICE_MAX_CONNECTIONS: 1 + NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0 + NCCL_ALGO: Ring + CUBLAS_WORKSPACE_CONFIG: :4096:8 + ENABLE_LIGHTWEIGHT_MODE: true +MODEL_ARGS: + --hidden-size: 512 + --num-attention-heads: 8 + --disable-bias-linear: true + --multi-latent-attention: true + --q-lora-rank: 192 + --kv-lora-rank: 64 + --qk-head-dim: 16 + --qk-pos-emb-head-dim: 8 + --v-head-dim: 16 + --hybrid-layer-pattern: WEC-|H-C-/W- + --spec: "[megatron.core.models.hybrid.hybrid_layer_specs hybrid_dsv4_stack_spec]" + --experimental-attention-variant: dsv4_hybrid + --dsa-indexer-n-heads: 64 + --dsa-indexer-head-dim: 128 + --dsa-indexer-topk: 512 + --dsa-indexer-loss-coeff: 0.01 + --dsa-indexer-use-sparse-loss: true + --csa-window-size: 128 + --csa-compress-ratios: ([0,4,128,4,0]) + --csa-compress-rotary-base: 40000 + --no-dsa-kernel-fusion: true + --num-experts: 8 + --expert-model-parallel-size: 1 + --expert-tensor-parallel-size: 1 + --moe-router-topk: 2 + --moe-token-dispatcher-type: alltoall + --moe-grouped-gemm: true + --moe-router-dtype: fp32 + --moe-router-load-balancing-type: aux_loss + --moe-aux-loss-coeff: 1.0e-4 + --enable-hyper-connections: true + --num-residual-streams: 4 + --mhc-sinkhorn-iterations: 20 + --mtp-num-layers: 1 + --mtp-loss-scaling-factor: 0.1 + --log-params-norm: true + --log-num-zeros-in-grad: true + --log-validation-ppl-to-tensorboard: true + --log-timers-to-tensorboard: true + --tensorboard-dir: ${TENSORBOARD_PATH} + --micro-batch-size: 4 + --global-batch-size: 32 + --seq-length: 1024 + --use-varlen-dataset: true + --mock-data: true + --varlen-mock-dataset-config-json: '\''{\"mode\":\"distribution\",\"type\":\"lognormal\",\"format\":\"thd\",\"min_seq_len\":1024,\"max_seq_len\":1024,\"mean_seq_len\":1024,\"lognormal_sigma\":1.1}\''' + --sequence-packing-scheduler: dp_balanced + --max-seqlen-per-dp-cp-rank: 512 + --pad-packed-seq-alignment: max + --no-pad-packed-seq-by-appending-dummy-seq: true + --calculate-per-token-loss: true + --thd-max-packed-sequences: 8 + --cp-partition-mode: contiguous + --dataloader-type: single + --tokenizer-type: NullTokenizer + --vocab-size: 100352 + --position-embedding-type: rope + --max-position-embeddings: 1024 + --train-iters: 50 + --timing-log-level: 0 + --lr-decay-iters: 320000 + --save: ${CHECKPOINT_SAVE_PATH} + --load: ${CHECKPOINT_LOAD_PATH} + --split: 949,50,1 + --distributed-backend: nccl + --lr: 0.00015 + --lr-decay-style: cosine + --min-lr: 1.0e-5 + --weight-decay: 1e-2 + --clip-grad: 1.0 + --lr-warmup-fraction: .01 + --log-interval: 1 + --save-interval: 25 + --eval-interval: 1000 + --eval-iters: 10 + --transformer-impl: transformer_engine + --tensor-model-parallel-size: 1 + --pipeline-model-parallel-size: 2 + --context-parallel-size: 2 + --sequence-parallel: true + --untie-embeddings-and-output-weights: true + --deterministic-mode: true + --no-gradient-accumulation-fusion: true + --attention-softmax-in-fp32: true + --use-mcore-models: true + --cuda-graph-impl: transformer_engine + --cuda-graph-dynamic-microbatches: true + --cuda-graph-modules: "[attn moe_router moe_preprocess]" + --te-rng-tracker: true + --ckpt-format: torch_dist + --data-cache-path: ${DATA_CACHE_PATH} + --bf16: true + --attention-backend: unfused + --log-memory-to-tensorboard: true +TEST_TYPE: ckpt-resume +METRICS: + - "lm loss" + - "num-zeros" + - "mtp_1 loss" diff --git a/tests/test_utils/recipes/h100/mamba.yaml b/tests/test_utils/recipes/h100/mamba.yaml index d0e0aba156a..95a4606dab0 100644 --- a/tests/test_utils/recipes/h100/mamba.yaml +++ b/tests/test_utils/recipes/h100/mamba.yaml @@ -81,6 +81,12 @@ products: # - environment: [lts] # disabled until triton is bumped # scope: [nightly] + - test_case: [hybrid_mr_mcore_te_tp1_pp2_dsv4_hybrid_mhc_mtp] + products: + - environment: [dev] + scope: [mr, mr-github, mr-github-slim] + platforms: [dgx_h100] + - test_case: [hybrid_mr_mcore_te_tp2_pp1_cp1_dgx_a100_1N8G] products: - environment: [dev] diff --git a/tests/unit_tests/data/test_get_batch.py b/tests/unit_tests/data/test_get_batch.py index 3d96caee968..13b3e82cc4b 100644 --- a/tests/unit_tests/data/test_get_batch.py +++ b/tests/unit_tests/data/test_get_batch.py @@ -15,7 +15,7 @@ ) from megatron.training.arguments import parse_args, validate_args from megatron.training.global_vars import destroy_global_vars, set_global_variables -from pretrain_hybrid import get_batch +from pretrain_hybrid import core_gpt_dataset_config_from_args, get_batch from tests.unit_tests.test_utilities import Utils @@ -70,6 +70,98 @@ def initialize_test_environment( return args +def test_sequence_packing_scheduler_forwards_process_groups(): + """Hybrid scheduler batching preserves THD metadata and model process groups.""" + args = MagicMock( + context_parallel_size=2, + sft=False, + dataloader_inter_document_masking=False, + create_attention_mask_in_dataloader=False, + hybrid_context_parallel=False, + sequence_packing_scheduler="dp_balanced", + ) + config = MagicMock( + virtual_pipeline_model_parallel_size=None, + pipeline_model_parallel_layout=None, + mtp_num_layers=1, + ) + data_iterator = object() + pg_collection = object() + tokens, labels, loss_mask, position_ids = (object() for _ in range(4)) + packed_seq_params, padding_mask = object(), object() + scheduler_batch = ( + tokens, + labels, + loss_mask, + None, + position_ids, + packed_seq_params, + padding_mask, + ) + + with ( + patch("pretrain_hybrid.get_args", return_value=args), + patch("pretrain_hybrid.core_transformer_config_from_args", return_value=config), + patch("pretrain_hybrid.mpu.get_tensor_model_parallel_rank", return_value=0), + patch("pretrain_hybrid.mtp_on_this_rank_func", return_value=True), + patch( + "pretrain_hybrid.get_batch_on_this_rank_for_sequence_packing", + return_value=scheduler_batch, + ) as scheduler_get_batch, + ): + batch = get_batch(data_iterator, vp_stage=0, pg_collection=pg_collection) + + assert batch == ( + None, + None, + None, + None, + labels, + None, + loss_mask, + None, + position_ids, + tokens, + padding_mask, + packed_seq_params, + ) + scheduler_get_batch.assert_called_once_with( + data_iterator, + vpp_size=None, + mtp_on_this_rank=True, + vp_stage=0, + pg_collection=pg_collection, + config=config, + ) + + +def test_varlen_dataset_config_uses_main_api(): + """Hybrid varlen data uses VarlenDatasetConfig with parsed mock settings.""" + args = MagicMock( + per_dataset_sequences_path=None, + use_varlen_dataset=True, + varlen_mock_dataset_config_json='{"mode": "distribution"}', + varlen_sbhd_validation=False, + ) + parsed_mock_config = {"mode": "distribution"} + expected_config = object() + + with ( + patch("pretrain_hybrid.build_tokenizer", return_value=object()), + patch("pretrain_hybrid.get_blend_and_blend_per_split", return_value=(None, None)), + patch("pretrain_hybrid.load_json_arg", return_value=parsed_mock_config) as load_json, + patch("pretrain_hybrid.VarlenDatasetConfig", return_value=expected_config) as config_class, + ): + config = core_gpt_dataset_config_from_args(args) + + assert config is expected_config + load_json.assert_called_once_with(args.varlen_mock_dataset_config_json) + config_kwargs = config_class.call_args.kwargs + assert config_kwargs["mock_dataset_config"] is parsed_mock_config + assert config_kwargs["sbhd_validation"] is False + assert "varlen_mock_dataset_config_json" not in config_kwargs + + def create_sft_data_iterator(max_seq_length: int = 1024): """Create a mock SFT data iterator matching the old SFTDataset output after DataLoader collation. @@ -190,8 +282,13 @@ def test_sft_batch(tp_size, pp_size, cp_size, seq_length): max_seqlen, position_ids, tokens, + padding_mask, + packed_seq_params, ) = get_batch(data_iterator) + assert padding_mask is None + assert packed_seq_params is None + is_first = mpu.is_pipeline_first_stage() is_last = mpu.is_pipeline_last_stage() seq_len_per_rank = seq_length // cp_size @@ -757,8 +854,13 @@ def test_pretrain_batch( max_seqlen, position_ids, tokens, + padding_mask, + packed_seq_params, ) = get_batch(data_iterator) + assert padding_mask is None + assert packed_seq_params is None + is_first = mpu.is_pipeline_first_stage() is_last = mpu.is_pipeline_last_stage() seq_len_per_rank = seq_length // cp_size @@ -986,8 +1088,13 @@ def test_hybrid_cp_batch(tp_size, cp_size, seq_length, create_attention_mask): max_seqlen, position_ids, tokens, + padding_mask, + packed_seq_params, ) = get_batch(data_iterator) + assert padding_mask is None + assert packed_seq_params is None + # Presence checks assert tokens is not None assert labels is not None diff --git a/tests/unit_tests/data/test_varlen_dataset.py b/tests/unit_tests/data/test_varlen_dataset.py new file mode 100644 index 00000000000..b3cc7631ecb --- /dev/null +++ b/tests/unit_tests/data/test_varlen_dataset.py @@ -0,0 +1,525 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Tests for the variable-length raw-text and mock pretraining datasets.""" + +import json +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from megatron.training.datasets.varlen_dataset import ( + MockVarlenDataset, + VarlenDataset, + VarlenDatasetConfig, + VarlenLowLevelDataset, + _looks_like_hf_id, + _raw_text_loader, +) + + +@pytest.mark.parametrize( + "path,expected", + [ + ("owner/dataset", True), + ("/tmp/data.jsonl", False), + ("./data.jsonl", False), + ("../data.jsonl", False), + ("data.jsonl", False), + ("", False), + (None, False), + ], +) +def test_looks_like_hf_id(path, expected): + assert _looks_like_hf_id(path) is expected + + +def _varlen_args(**overrides): + values = { + "use_varlen_dataset": False, + "varlen_sbhd_validation": False, + "varlen_mock_dataset_config_json": None, + "sft": False, + "fim_data": False, + "mock_data": False, + "hybrid_context_parallel": False, + "num_experts": None, + "reset_position_ids": False, + "reset_attention_mask": False, + "dataloader_inter_document_masking": False, + "create_attention_mask_in_dataloader": True, + "sequence_packing_scheduler": None, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def test_varlen_validation_selects_default_scheduler(): + from megatron.training.arguments import _validate_varlen_dataset_args + + args = _varlen_args(use_varlen_dataset=True) + _validate_varlen_dataset_args(args) + + assert args.sequence_packing_scheduler == "dp_balanced" + assert args.create_attention_mask_in_dataloader is False + + +def test_varlen_validation_preserves_explicit_scheduler(): + from megatron.training.arguments import _validate_varlen_dataset_args + + args = _varlen_args(use_varlen_dataset=True, sequence_packing_scheduler="dp_balanced") + _validate_varlen_dataset_args(args) + + assert args.sequence_packing_scheduler == "dp_balanced" + + +@pytest.mark.parametrize( + "overrides", + [ + {"varlen_sbhd_validation": True}, + {"varlen_mock_dataset_config_json": "{}", "mock_data": True}, + {"use_varlen_dataset": True, "varlen_mock_dataset_config_json": "{}"}, + {"use_varlen_dataset": True, "sft": True}, + {"use_varlen_dataset": True, "fim_data": True}, + {"use_varlen_dataset": True, "reset_position_ids": True}, + {"use_varlen_dataset": True, "reset_attention_mask": True}, + {"use_varlen_dataset": True, "dataloader_inter_document_masking": True}, + { + "use_varlen_dataset": True, + "varlen_sbhd_validation": True, + "sequence_packing_scheduler": "dp_balanced", + }, + {"use_varlen_dataset": True, "varlen_sbhd_validation": True, "mock_data": True}, + {"use_varlen_dataset": True, "hybrid_context_parallel": True}, + {"use_varlen_dataset": True, "varlen_sbhd_validation": True, "num_experts": 8}, + ], +) +def test_varlen_validation_rejects_incompatible_options(overrides): + from megatron.training.arguments import _validate_varlen_dataset_args + + with pytest.raises(AssertionError): + _validate_varlen_dataset_args(_varlen_args(**overrides)) + + +def test_raw_text_loader(): + assert _raw_text_loader({"text": "hello"}) == "hello" + assert _raw_text_loader({"text": None}) == "" + with pytest.raises(ValueError, match="must be a string"): + _raw_text_loader({"text": [1, 2]}) + + +def _write_jsonl(tmp_path: Path, rows): + path = tmp_path / "data.jsonl" + path.write_text("\n".join(json.dumps(row) for row in rows)) + return str(path) + + +def test_low_level_loads_json_array_raw_text(tmp_path): + pytest.importorskip("datasets") + pytest.importorskip("pandas") + path = tmp_path / "data.json" + path.write_text(json.dumps([{"text": "one"}, {"text": "two", "id": 2}])) + + dataset = VarlenLowLevelDataset(str(path)) + + assert dataset.schema_name == "pretrain-text" + assert [dataset[0], dataset[1]] == ["one", "two"] + + +def test_low_level_loads_jsonl_raw_text(tmp_path): + pytest.importorskip("datasets") + pytest.importorskip("pandas") + path = _write_jsonl( + tmp_path, + [ + {"text": "document one", "url": "https://example/1"}, + {"text": "document two", "url": "https://example/2"}, + ], + ) + + dataset = VarlenLowLevelDataset(path) + + assert len(dataset) == 2 + assert dataset[1] == "document two" + + +def test_low_level_rejects_non_text_schema(tmp_path): + pytest.importorskip("datasets") + pytest.importorskip("pandas") + path = _write_jsonl(tmp_path, [{"messages": [{"role": "user", "content": "hello"}]}]) + + with pytest.raises(ValueError, match="requires a raw pretraining 'text' column"): + VarlenLowLevelDataset(path) + + +class _FakeTokenizer: + def __init__(self, eod: int = 0, pad=None): + self._eod = eod + self._pad = pad + + @property + def eod(self): + return self._eod + + @property + def pad(self): + return self._pad + + @property + def vocab_size(self): + return 128 + + def tokenize(self, text): + return [ord(char) % 100 + 1 for char in text] + + +def test_varlen_dataset_config_owns_mock_and_sbhd_fields(): + config_args = { + "random_seed": 123, + "sequence_length": 8, + "tokenizer": _FakeTokenizer(), + "reset_position_ids": False, + "reset_attention_mask": False, + "eod_mask_loss": False, + "create_attention_mask": False, + "context_parallel_size": 1, + } + mock_config = {"mode": "file", "path": "lengths.csv"} + + config = VarlenDatasetConfig( + **config_args, mock_dataset_config=mock_config, sbhd_validation=True + ) + + assert config.mock_dataset_config == mock_config + assert config.sbhd_validation is True + with pytest.raises(AssertionError, match="hybrid context parallelism"): + VarlenDatasetConfig(**config_args, hybrid_context_parallel=True) + + +def _make_config(tokenizer, seq_length=64, *, cp=1, dp=1, sp=1, sbhd=False, eod_mask=False): + return SimpleNamespace( + tokenizer=tokenizer, + sequence_length=seq_length, + reset_position_ids=False, + create_attention_mask=False, + reset_attention_mask=False, + eod_mask_loss=eod_mask, + sbhd_validation=sbhd, + data_parallel_size=dp, + context_parallel_size=cp, + sequence_parallel_size=sp, + hybrid_context_parallel=False, + ) + + +def _make_varlen(items, config): + dataset = VarlenDataset.__new__(VarlenDataset) + dataset.config = config + dataset.dataset = items + dataset.indices = np.arange(len(items)) + return dataset + + +def _make_mock_varlen(token_arrays, config): + dataset = MockVarlenDataset.__new__(MockVarlenDataset) + dataset.config = config + dataset.dataset = token_arrays + dataset.indices = np.arange(len(token_arrays)) + return dataset + + +def test_getitem_thd_raw_text_keys_and_shapes(): + dataset = _make_varlen( + ["hello world"], _make_config(_FakeTokenizer(eod=0, pad=7)) + ) + + sample = dataset[0] + + assert set(sample) == { + "tokens", + "labels", + "loss_mask", + "position_ids", + "original_seq_len", + "padded_seq_len", + } + size = sample["tokens"].numel() + assert sample["labels"].numel() == size + assert sample["loss_mask"].numel() == size + assert sample["position_ids"].numel() == size + assert sample["padded_seq_len"].item() == size + + +def test_getitem_thd_pad_mask_keeps_real_eod(): + tokenizer = _FakeTokenizer(eod=0, pad=None) + dataset = _make_varlen(["abc"], _make_config(tokenizer, cp=2)) + + sample = dataset[0] + + assert sample["labels"].tolist()[2] == tokenizer.eod + assert sample["loss_mask"].tolist() == [1.0, 1.0, 1.0, 0.0] + + +def test_getitem_thd_pads_to_cp_divisor(): + dataset = _make_varlen( + ["abcde"], _make_config(_FakeTokenizer(eod=0, pad=7), cp=2) + ) + + sample = dataset[0] + + assert sample["padded_seq_len"].item() % 4 == 0 + + +def test_getitem_empty_text_is_nonempty(): + dataset = _make_varlen([""], _make_config(_FakeTokenizer(eod=0, pad=7))) + + sample = dataset[0] + + assert sample["tokens"].numel() >= 1 + assert sample["labels"].numel() == sample["tokens"].numel() + + +def test_getitem_exact_capacity_reserves_eod(): + tokenizer = _FakeTokenizer(eod=0, pad=7) + dataset = _make_varlen(["abcde"], _make_config(tokenizer, seq_length=4)) + + sample = dataset[0] + + assert sample["tokens"].numel() == 4 + assert sample["labels"][-1].item() == tokenizer.eod + assert sample["loss_mask"][-1].item() == 1.0 + + +@pytest.mark.parametrize("sbhd", [False, True]) +def test_getitem_honors_eod_mask_loss(sbhd): + tokenizer = _FakeTokenizer(eod=0, pad=7) + tokenizer.tokenize = lambda _: [1, tokenizer.eod, 2] + dataset = _make_varlen( + ["abc"], _make_config(tokenizer, seq_length=8, sbhd=sbhd, eod_mask=True) + ) + + sample = dataset[0] + + assert sample["tokens"][1].item() == tokenizer.eod + assert sample["loss_mask"][1].item() == 0.0 + assert sample["labels"][0].item() == tokenizer.eod + assert sample["loss_mask"][0].item() == 1.0 + + +def test_getitem_sbhd_pads_to_sequence_length(): + dataset = _make_varlen( + ["abc"], _make_config(_FakeTokenizer(eod=0, pad=None), seq_length=8, sbhd=True) + ) + + sample = dataset[0] + + assert set(sample) == {"tokens", "labels", "loss_mask", "position_ids"} + assert sample["tokens"].numel() == 8 + assert sample["loss_mask"].tolist() == [1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0] + + +def test_mock_getitem_thd_keys_and_pad_fallback(): + dataset = _make_mock_varlen( + [np.array([1, 2, 3, 4], dtype=np.int64)], + _make_config(_FakeTokenizer(eod=0, pad=None), cp=2), + ) + + sample = dataset[0] + + assert set(sample) == { + "tokens", + "labels", + "loss_mask", + "position_ids", + "original_seq_len", + "padded_seq_len", + } + assert sample["padded_seq_len"].item() % 4 == 0 + + +def test_mock_getitem_truncates_to_sequence_length(): + dataset = _make_mock_varlen( + [np.arange(1, 10, dtype=np.int64)], + _make_config(_FakeTokenizer(eod=0, pad=7), seq_length=4), + ) + + sample = dataset[0] + + assert sample["original_seq_len"].item() == 4 + assert sample["tokens"].numel() == 4 + + +def test_mock_getitem_length_one_is_nonempty(): + dataset = _make_mock_varlen( + [np.array([], dtype=np.int64)], _make_config(_FakeTokenizer(eod=0, pad=7)) + ) + + sample = dataset[0] + + assert sample["original_seq_len"].item() == 1 + assert sample["tokens"].numel() == 1 + + +def test_unpack_batch_normalizes_varlen_samples(): + from megatron.core.datasets.data_schedule_utils import _unpack_batch + + batch = [ + { + "tokens": torch.arange(4, dtype=torch.int64).view(1, 4), + "labels": torch.arange(4, dtype=torch.int64).view(1, 4), + "loss_mask": torch.ones(1, 4), + "position_ids": torch.arange(4, dtype=torch.int64).view(1, 4), + "padded_seq_len": torch.tensor([4], dtype=torch.int32), + } + ] + + output = _unpack_batch(batch) + + assert output[0]["tokens"].shape == (4,) + assert output[0]["original_seq_len"].item() == 4 + + +def test_unpack_batch_requires_varlen_lengths(): + from megatron.core.datasets.data_schedule_utils import _unpack_batch + + with pytest.raises(KeyError, match="padded_seq_len"): + _unpack_batch([{"cu_seqlens": torch.tensor([0, 4], dtype=torch.int32)}]) + + +def _build_varlen_for_loader(items, config, num_samples): + from megatron.core.datasets.utils import Split + + dataset = VarlenDataset.__new__(VarlenDataset) + dataset.config = config + dataset.dataset = items + dataset.indices = np.arange(len(items)) + dataset.num_samples = num_samples + dataset.index_split = Split.train + return dataset + + +def _loader_args(*, use_varlen, sbhd, scheduler, mbs, gbs=None): + return SimpleNamespace( + dataloader_type="single", + micro_batch_size=mbs, + global_batch_size=mbs if gbs is None else gbs, + full_validation=False, + num_workers=0, + use_varlen_dataset=use_varlen, + varlen_sbhd_validation=sbhd, + sequence_packing_scheduler=scheduler, + ) + + +def test_sbhd_validation_dataloader_uses_default_collate(): + from megatron.core import parallel_state + from megatron.training.datasets.data_samplers import build_pretraining_data_loader + from megatron.training.global_vars import destroy_global_vars, set_args + from tests.unit_tests.test_utilities import Utils + + Utils.initialize_model_parallel(1, 1) + try: + tokenizer = _FakeTokenizer(eod=0, pad=7) + seq_len, micro_batch_size = 16, 2 + data_parallel_size = parallel_state.get_data_parallel_world_size() + num_samples = micro_batch_size * data_parallel_size * 4 + dataset = _build_varlen_for_loader( + ["hello world"] * num_samples, + _make_config(tokenizer, seq_length=seq_len, sbhd=True), + num_samples, + ) + set_args( + _loader_args( + use_varlen=True, + sbhd=True, + scheduler=None, + mbs=micro_batch_size, + ) + ) + + batch = next(iter(build_pretraining_data_loader(dataset, consumed_samples=0))) + + assert isinstance(batch, dict) + assert batch["tokens"].shape == (micro_batch_size, seq_len) + finally: + destroy_global_vars() + Utils.destroy_model_parallel() + + +def test_thd_dataloader_uses_identity_collate(): + from megatron.core import parallel_state + from megatron.training.datasets.data_samplers import build_pretraining_data_loader + from megatron.training.global_vars import destroy_global_vars, set_args + from tests.unit_tests.test_utilities import Utils + + Utils.initialize_model_parallel(1, 1) + try: + tokenizer = _FakeTokenizer(eod=0, pad=7) + micro_batch_size = 2 + data_parallel_size = parallel_state.get_data_parallel_world_size() + num_samples = micro_batch_size * data_parallel_size * 4 + values = ["a", "abcdef", "xy", "qwerty"] + dataset = _build_varlen_for_loader( + [values[index % len(values)] for index in range(num_samples)], + _make_config(tokenizer), + num_samples, + ) + set_args( + _loader_args( + use_varlen=True, + sbhd=False, + scheduler="dp_balanced", + mbs=micro_batch_size, + ) + ) + + batch = next(iter(build_pretraining_data_loader(dataset, consumed_samples=0))) + + assert isinstance(batch, list) + assert len(batch) == micro_batch_size + assert "padded_seq_len" in batch[0] + finally: + destroy_global_vars() + Utils.destroy_model_parallel() + + +def test_packing_scheduler_dataloader_yields_microbatches(): + from megatron.core import parallel_state + from megatron.training.datasets.data_samplers import build_pretraining_data_loader + from megatron.training.global_vars import destroy_global_vars, set_args + from tests.unit_tests.test_utilities import Utils + + Utils.initialize_model_parallel(1, 1) + try: + tokenizer = _FakeTokenizer(eod=0, pad=7) + micro_batch_size = 2 + num_microbatches = 3 + data_parallel_size = parallel_state.get_data_parallel_world_size() + global_batch_size = micro_batch_size * data_parallel_size * num_microbatches + num_samples = global_batch_size * 2 + values = ["a", "abcdef", "xy", "qwerty"] + dataset = _build_varlen_for_loader( + [values[index % len(values)] for index in range(num_samples)], + _make_config(tokenizer, dp=data_parallel_size), + num_samples, + ) + set_args( + _loader_args( + use_varlen=True, + sbhd=False, + scheduler="dp_balanced", + mbs=micro_batch_size, + gbs=global_batch_size, + ) + ) + + batch = next(iter(build_pretraining_data_loader(dataset, consumed_samples=0))) + + assert isinstance(batch, list) + assert len(batch) == micro_batch_size + assert "padded_seq_len" in batch[0] + finally: + destroy_global_vars() + Utils.destroy_model_parallel() diff --git a/tests/unit_tests/dist_checkpointing/test_optimizer.py b/tests/unit_tests/dist_checkpointing/test_optimizer.py index f93e09a43b7..fb0f45f627f 100644 --- a/tests/unit_tests/dist_checkpointing/test_optimizer.py +++ b/tests/unit_tests/dist_checkpointing/test_optimizer.py @@ -79,6 +79,25 @@ def sharded_state_dict(self): return sharded_state_dict +class NativeFp32Model(torch.nn.Module): + """Three parameters that can be converted to an interleaved BF16/FP32/BF16 group.""" + + def __init__(self): + super().__init__() + self.pre = torch.nn.Linear(8, 8, bias=False) + self.gate = torch.nn.Parameter(torch.zeros(24, dtype=torch.float32)) + self.post = torch.nn.Linear(8, 8, bias=False) + self.config = TransformerConfig( + hidden_size=8, num_attention_heads=1, num_layers=1, bf16=True + ) + + def sharded_state_dict(self): + return { + key: ShardedTensor.from_rank_offsets(key, value) + for key, value in self.state_dict(keep_vars=True).items() + } + + class SwigluFactoryModel(torch.nn.Module): def __init__(self, pp_separate_model: bool = False): super().__init__() @@ -238,6 +257,59 @@ def test_optimizer_params(self, tmp_path_dist_ckpt): ] ) + def test_float16_optimizer_with_native_fp32_params(self): + """Native FP32 state ids must remain correct between two BF16 parameters.""" + from megatron.core.optimizer import OptimizerConfig + from megatron.core.optimizer.optimizer import Float16OptimizerWithFloat16Params + from megatron.core.transformer.module import ( + convert_module_to_dtype_except_fp32_marked, + mark_keep_in_fp32, + ) + + Utils.initialize_model_parallel(1, 1) + model = NativeFp32Model().cuda() + model.gate = mark_keep_in_fp32(model.gate) + convert_module_to_dtype_except_fp32_marked(model, torch.bfloat16) + assert model.pre.weight.dtype == torch.bfloat16 + assert model.gate.dtype == torch.float32 + assert model.post.weight.dtype == torch.bfloat16 + + # Use an explicit BF16/FP32/BF16 optimizer order. Module.parameters() + # would yield the root gate before parameters owned by child modules. + ordered_params = [model.pre.weight, model.gate, model.post.weight] + for param in ordered_params: + param.grad = torch.zeros_like(param) + inner_optim = Adam(ordered_params) + inner_optim.step() + + optim = Float16OptimizerWithFloat16Params( + inner_optim, + OptimizerConfig(optimizer='adam', lr=1e-4, bf16=True), + None, + lambda opt, cfg: None, + ) + sharded_state_dict = optim.sharded_state_dict(model.sharded_state_dict()) + + # FP32 main copies pair with the BF16 params only, in optimizer order. + fp32_params = sharded_state_dict['fp32_from_fp16_params'][0] + assert [(sharded.key, tuple(sharded.data.shape)) for sharded in fp32_params] == [ + ('optimizer.state.fp32_param.pre.weight', (8, 8)), + ('optimizer.state.fp32_param.post.weight', (8, 8)), + ] + + # Per-param state maps every param, including the native FP32 one, to the right key. + state = sharded_state_dict['optimizer']['state'] + expected = {0: ('pre.weight', (8, 8)), 1: ('gate', (24,)), 2: ('post.weight', (8, 8))} + for param_id, (model_key, shape) in expected.items(): + for state_key in ('exp_avg', 'exp_avg_sq'): + sharded = state[param_id][state_key] + assert sharded.key == f'optimizer.state.{state_key}.{model_key}', sharded.key + assert tuple(sharded.data.shape) == shape, ( + param_id, + sharded.key, + sharded.data.shape, + ) + def initialize_pp_agnostic_model(pre_process=True, post_process=True, seed=0, **config_kwargs): torch.manual_seed(seed) diff --git a/tests/unit_tests/distributed/test_finalize_model_grads.py b/tests/unit_tests/distributed/test_finalize_model_grads.py index 80d143a89a3..6c16f87918b 100644 --- a/tests/unit_tests/distributed/test_finalize_model_grads.py +++ b/tests/unit_tests/distributed/test_finalize_model_grads.py @@ -1,6 +1,7 @@ # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. import inspect import os +from types import SimpleNamespace import pytest import torch @@ -11,6 +12,7 @@ from megatron.core.distributed.finalize_model_grads import ( _allreduce_non_tensor_model_parallel_grads, _allreduce_word_embedding_grads, + _update_router_expert_bias, _update_router_qb_beta, finalize_model_grads, reset_model_temporary_tensors, @@ -44,6 +46,30 @@ def finish_grad_sync(self, force_all_reduce=False): self.finish_grad_sync_calls += 1 +class _HashRouterWithoutExpertBias(torch.nn.Module): + """Match hash-router layers, which intentionally do not own expert-bias state.""" + + def __init__(self): + super().__init__() + self.expert_bias = None + self.local_tokens_per_expert = None + + +def test_hash_router_without_expert_bias_is_ignored(): + router = _HashRouterWithoutExpertBias() + config = SimpleNamespace( + moe_router_enable_expert_bias=True, + moe_router_load_balancing_type="none", + moe_router_bias_update_rate=0.25, + ) + + reset_model_temporary_tensors(config, [router]) + _update_router_expert_bias([router], config) + + assert router.expert_bias is None + assert router.local_tokens_per_expert is None + + def _router_expert_bias_config(): return TransformerConfig( num_layers=1, diff --git a/tests/unit_tests/fusions/test_bias_dropout_fusion.py b/tests/unit_tests/fusions/test_bias_dropout_fusion.py index f8b23900543..a7c63626e93 100644 --- a/tests/unit_tests/fusions/test_bias_dropout_fusion.py +++ b/tests/unit_tests/fusions/test_bias_dropout_fusion.py @@ -319,3 +319,86 @@ def test_fp32_residual_precision_advantage(self): f"fp32 residual error ({err_fp32:.6e}) should be less than " f"bf16 residual error ({err_bf16:.6e})" ) + + +# ============================================================================ +# Tests for the mHC recompute path of get_bias_dropout_add +# ============================================================================ +# +# When ``mhc_recompute_manager`` is provided, ``get_bias_dropout_add`` returns +# a closure that wraps the underlying BDA in ``CheckpointWithoutOutput`` and +# auto-registers with the supplied ``CheckpointWithoutOutputManager``. These tests cover +# that branch (which is otherwise only invoked indirectly from the mHC layer +# forward path). + + +class TestBiasDropoutAddMhcRecompute: + """Direct coverage for ``_get_checkpointed_bda``.""" + + def setup_method(self, method): + from megatron.core.tensor_parallel.random import initialize_rng_tracker + from tests.unit_tests.test_utilities import Utils + + Utils.initialize_model_parallel() + initialize_rng_tracker(force_reset=True) + + def teardown_method(self, method): + from tests.unit_tests.test_utilities import Utils + + Utils.destroy_model_parallel() + + @pytest.mark.parametrize("fused", [False, True]) + @pytest.mark.parametrize("with_bias", [True, False]) + def test_checkpointed_bda_forward_backward(self, fused, with_bias): + """Closure runs forward+backward and registers with the manager.""" + from megatron.core.tensor_parallel.random import CheckpointWithoutOutputManager + + torch.manual_seed(0) + manager = CheckpointWithoutOutputManager() + bda = get_bias_dropout_add(training=True, fused=fused, mhc_recompute_manager=manager) + + x = torch.randn(8, 4, 16, device="cuda", requires_grad=True) + residual = torch.randn_like(x, requires_grad=True) + bias = torch.zeros(16, device="cuda") if with_bias else None + x_with_bias = (x, bias) if with_bias else x + + out = bda(x_with_bias, residual, 0.0) + assert out.shape == x.shape + assert out.dtype == x.dtype + assert len(manager.checkpoints) == 1, "checkpoint should auto-register with manager" + + loss = out.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss) + loss.backward() + + assert x.grad is not None and torch.isfinite(x.grad).all() + assert residual.grad is not None and torch.isfinite(residual.grad).all() + + def test_checkpointed_bda_chained_managers(self): + """Two checkpointed BDAs chained on one manager both register.""" + from megatron.core.tensor_parallel.random import CheckpointWithoutOutputManager + + torch.manual_seed(0) + manager = CheckpointWithoutOutputManager() + bda = get_bias_dropout_add(training=True, fused=False, mhc_recompute_manager=manager) + + x = torch.randn(4, 2, 8, device="cuda", requires_grad=True) + residual = torch.randn_like(x, requires_grad=True) + + y1 = bda((x, None), residual, 0.0) + y2 = bda((y1, None), residual, 0.0) + + assert len(manager.checkpoints) == 2, "each call should register a new checkpoint" + loss = y2.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss) + loss.backward() + assert x.grad is not None + + def test_get_bda_without_manager_unchanged(self): + """The default (manager=None) path returns the regular BDA, not a closure.""" + unfused = get_bias_dropout_add(training=True, fused=False) + fused = get_bias_dropout_add(training=False, fused=True) + # Both must be callable; neither should be the mHC closure (which has __closure__ over manager). + assert callable(unfused) and callable(fused) + assert getattr(unfused, "__name__", "") != "_checkpointed_bda" + assert getattr(fused, "__name__", "") != "_checkpointed_bda" diff --git a/tests/unit_tests/fusions/test_fused_mhc_kernels.py b/tests/unit_tests/fusions/test_fused_mhc_kernels.py new file mode 100644 index 00000000000..d05870852f4 --- /dev/null +++ b/tests/unit_tests/fusions/test_fused_mhc_kernels.py @@ -0,0 +1,1427 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Unit tests for unified fused mHC kernels and native implementations. + +Each test compares the fused kernel's forward output AND backward gradients +against a pure-PyTorch differentiable reference to catch numerical drift +introduced by kernel fusion. The public fused API tests also exercise backend +dispatch and fallback selection through the same entry points used by mHC. +""" + +import math +from typing import Optional + +import pytest +import torch +from torch import Tensor + +from megatron.core.fusions.fused_mhc_kernels import is_cutile_available, is_triton_available +from megatron.core.transformer.hyper_connection import ( + native_h_aggregate, + native_h_post_bda, + native_proj_rms, + native_sinkhorn, +) + +_require_cutile = pytest.mark.skipif( + not is_cutile_available(), reason="cuTile unavailable for current device" +) +_require_triton = pytest.mark.skipif(not is_triton_available(), reason="Triton not installed") + + +@pytest.fixture(autouse=True) +def _skip_without_cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + +DTYPE = torch.bfloat16 +DEVICE = "cuda" +FWD_ATOL, FWD_RTOL = 2e-2, 2e-2 +BWD_ATOL, BWD_RTOL = 5e-2, 5e-2 +RAND_LO, RAND_HI = -0.1, 0.1 +COSINE_SIM_THRESH = 0.999 + + +def _assert_cosine_similar(a: Tensor, b: Tensor, threshold: float, msg: str = ""): + """Assert that flattened tensors have cosine similarity >= threshold.""" + a_flat = a.flatten().float() + b_flat = b.flatten().float() + sim = torch.nn.functional.cosine_similarity(a_flat.unsqueeze(0), b_flat.unsqueeze(0)).item() + assert sim >= threshold, ( + f"{msg}: cosine similarity {sim:.6f} < {threshold} " + f"(max_abs_diff={torch.max(torch.abs(a_flat - b_flat)):.6e})" + ) + + +def _rand(*shape, **kwargs): + """Uniform in [RAND_LO, RAND_HI] to keep magnitudes small for bf16 stability.""" + return torch.empty(*shape, dtype=DTYPE, device=DEVICE, **kwargs).uniform_(RAND_LO, RAND_HI) + + +def _info(): + if is_triton_available() and is_cutile_available(): + backend = "triton+cuTile" + elif is_triton_available(): + backend = "triton+native" + elif is_cutile_available(): + backend = "cuTile" + else: + backend = "native" + print(f"\n [backend: {backend}]") + + +# ============================================================================ +# Pure-PyTorch differentiable references (used by both fwd AND bwd tests) +# ============================================================================ + + +def _ref_sinkhorn(logits: Tensor, num_iters: int, eps: float = 1e-6) -> Tensor: + M = logits.softmax(dim=-1) + eps + M = M / (M.sum(dim=-2, keepdim=True) + eps) + for _ in range(num_iters - 1): + M = M / (M.sum(dim=-1, keepdim=True) + eps) + M = M / (M.sum(dim=-2, keepdim=True) + eps) + return M + + +def _ref_h_aggregate(x: Tensor, h_pre: Tensor) -> Tensor: + return (x * h_pre.unsqueeze(-1)).sum(dim=2) + + +def _ref_h_post_bda( + h_res: Tensor, orig_res: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] +) -> Tensor: + s, b, n, C = orig_res.shape + h_res_batched = h_res.view(s * b, n, n) + orig_batched = orig_res.view(s * b, n, C) + mixed = torch.bmm(h_res_batched.transpose(1, 2), orig_batched).view(s, b, n, C) + x_exp = h_post.unsqueeze(-1) * x.unsqueeze(2) + out = x_exp + mixed + if bias is not None: + out = out + h_post.unsqueeze(-1) * bias.view(1, 1, 1, C) + return out + + +def _h_post_bda_transpose_case(): + h_res = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]], dtype=DTYPE, device=DEVICE) + orig = torch.tensor([[[[10.0, 100.0], [1.0, 2.0]]]], dtype=DTYPE, device=DEVICE) + h_post = torch.zeros(1, 1, 2, dtype=DTYPE, device=DEVICE) + x = torch.zeros(1, 1, 2, dtype=DTYPE, device=DEVICE) + expected = torch.tensor([[[[13.0, 106.0], [24.0, 208.0]]]], dtype=DTYPE, device=DEVICE) + return h_res, orig, h_post, x, expected + + +def _ref_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6): + proj = torch.matmul(x, weight.t()) + norm = x.norm(dim=-1, keepdim=True) + K = x.shape[-1] + r = 1.0 / (norm / math.sqrt(K) + eps) + return proj, r + + +def _ref_proj_rms_compute_h( + x: Tensor, + weight: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + bias: Tensor, + n: int, + eps: float = 1e-6, + compute_h_eps: float = 1e-6, +): + """Reference: fused proj_rms + compute_h.""" + proj = torch.matmul(x, weight.t()) + norm = x.norm(dim=-1, keepdim=True) + K = x.shape[-1] + r = norm / math.sqrt(K) # [M, 1] + N = proj.shape[-1] + alpha = torch.cat([alpha_pre.expand(n), alpha_post.expand(n), alpha_res.expand(N - 2 * n)]) + h = proj * alpha.unsqueeze(0) / (r + eps) + bias.unsqueeze(0) + h_pre = h[..., :n].sigmoid() + compute_h_eps + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res = h[..., 2 * n :] + return h_pre, h_post, h_res, r + + +# ============================================================================ +# Sinkhorn +# ============================================================================ + + +class TestNativeSinkhorn: + """Tests for the native SinkhornKnopp implementation.""" + + @pytest.mark.parametrize("s,b,n,iters", [(2, 4, 4, 5), (1, 1, 2, 10)]) + def test_fwd_bwd_vs_torch_reference(self, s, b, n, iters): + """native_sinkhorn fwd output and bwd grad must match the inline PyTorch reference.""" + _info() + eps = 1e-6 + data = _rand(s, b, n, n) + grad_out = _rand(s, b, n, n) + + # -- native_sinkhorn path (autograd.Function) -- + inp_f = data.clone().requires_grad_(True) + out_f = native_sinkhorn(inp_f, iters, eps) + out_f.backward(grad_out) + grad_f = inp_f.grad.clone() + + # -- inline torch reference (fully differentiable) -- + inp_r = data.clone().requires_grad_(True) + out_r = _ref_sinkhorn(inp_r, iters, eps) + out_r.backward(grad_out) + grad_r = inp_r.grad.clone() + + torch.testing.assert_close(out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(grad_f, grad_r, atol=BWD_ATOL, rtol=BWD_RTOL) + + +class TestFusedSinkhorn: + """Public fused sinkhorn dispatch/fallback plus numerical correctness.""" + + @pytest.mark.flaky_in_dev + @_require_cutile + @pytest.mark.parametrize("s,b,n,iters", [(2, 4, 4, 5), (1, 1, 2, 10)]) + def test_fwd_bwd_vs_reference(self, s, b, n, iters): + """E2E: public fused fwd output and bwd grad must match the PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import fused_sinkhorn + + _info() + eps = 1e-6 + data = _rand(s, b, n, n) + grad_out = _rand(s, b, n, n) + + # -- fused path -- + inp_f = data.clone().requires_grad_(True) + out_f = fused_sinkhorn(inp_f, iters, eps) + out_f.backward(grad_out) + grad_f = inp_f.grad.clone() + + # -- reference path (fully differentiable) -- + inp_r = data.clone().requires_grad_(True) + out_r = _ref_sinkhorn(inp_r, iters, eps) + out_r.backward(grad_out) + grad_r = inp_r.grad.clone() + + torch.testing.assert_close(out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(grad_f, grad_r, atol=BWD_ATOL, rtol=BWD_RTOL) + + +# ============================================================================ +# H_aggregate +# ============================================================================ + + +class TestNativeHAggregate: + """Tests for native_h_aggregate.""" + + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 1, 2, 256)]) + def test_fwd_bwd_vs_torch_reference(self, s, b, n, C): + _info() + x_data = _rand(s, b, n, C) + h_data = _rand(s, b, n) + grad_out = _rand(s, b, C) + + xf = x_data.clone().requires_grad_(True) + hf = h_data.clone().requires_grad_(True) + of = native_h_aggregate(xf, hf) + of.backward(grad_out) + + xr = x_data.clone().requires_grad_(True) + hr = h_data.clone().requires_grad_(True) + oref = _ref_h_aggregate(xr, hr) + oref.backward(grad_out) + + torch.testing.assert_close(of, oref, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(xf.grad, xr.grad, atol=BWD_ATOL, rtol=BWD_RTOL) + torch.testing.assert_close(hf.grad, hr.grad, atol=BWD_ATOL, rtol=BWD_RTOL) + + +class TestFusedHAggregate: + """Public fused h_aggregate dispatch/fallback plus numerical correctness.""" + + @pytest.mark.flaky_in_dev + @_require_cutile + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 1, 2, 256)]) + def test_fwd_bwd_vs_reference(self, s, b, n, C): + """E2E: public fused fwd output and bwd grads must match the PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import fused_h_aggregate + + _info() + x_data = _rand(s, b, n, C) + h_data = _rand(s, b, n) + grad_out = _rand(s, b, C) + + # -- fused path -- + xf = x_data.clone().requires_grad_(True) + hf = h_data.clone().requires_grad_(True) + of = fused_h_aggregate(xf, hf) + of.backward(grad_out) + + # -- reference path -- + xr = x_data.clone().requires_grad_(True) + hr = h_data.clone().requires_grad_(True) + oref = _ref_h_aggregate(xr, hr) + oref.backward(grad_out) + + torch.testing.assert_close(of, oref, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(xf.grad, xr.grad, atol=BWD_ATOL, rtol=BWD_RTOL) + torch.testing.assert_close(hf.grad, hr.grad, atol=BWD_ATOL, rtol=BWD_RTOL) + + +class TestTritonHAggregate: + """Tests for Triton h_aggregate forward against PyTorch reference.""" + + @_require_triton + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 1, 2, 256), (64, 8, 4, 4096)]) + def test_fwd_vs_reference(self, s, b, n, C): + from megatron.core.fusions.fused_mhc_kernels import _triton_h_aggregate_fwd + + _info() + x_data = _rand(s, b, n, C) + h_data = _rand(s, b, n) + + out_t = _triton_h_aggregate_fwd(x_data, h_data) + out_r = _ref_h_aggregate(x_data, h_data) + + torch.testing.assert_close(out_t, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + + +# ============================================================================ +# H_post BDA +# ============================================================================ + + +class TestNativeHPostBDA: + """Tests for native_h_post_bda.""" + + def test_forward_uses_h_res_transpose(self): + h_res, orig, h_post, x, expected = _h_post_bda_transpose_case() + out = native_h_post_bda(h_res, orig, h_post, x, bias=None) + + torch.testing.assert_close(out, expected, atol=0.0, rtol=0.0) + + @pytest.mark.parametrize("with_bias", [True, False]) + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256)]) + def test_fwd_bwd_vs_torch_reference(self, s, b, n, C, with_bias): + _info() + hr_data = _rand(s, b, n, n) + orig_data = _rand(s, b, n, C) + hp_data = _rand(s, b, n) + x_data = _rand(s, b, C) + bias_data = _rand(C) if with_bias else None + grad_out = _rand(s, b, n, C) + + def _make_inputs(): + hr = hr_data.clone().requires_grad_(True) + orig = orig_data.clone().requires_grad_(True) + hp = hp_data.clone().requires_grad_(True) + x = x_data.clone().requires_grad_(True) + bi = bias_data.clone().requires_grad_(True) if with_bias else None + return hr, orig, hp, x, bi + + hr_f, orig_f, hp_f, x_f, bi_f = _make_inputs() + out_f = native_h_post_bda(hr_f, orig_f, hp_f, x_f, bi_f) + out_f.backward(grad_out) + + hr_r, orig_r, hp_r, x_r, bi_r = _make_inputs() + out_r = _ref_h_post_bda(hr_r, orig_r, hp_r, x_r, bi_r) + out_r.backward(grad_out) + + torch.testing.assert_close(out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + for name, gf, gr in [ + ("h_res", hr_f.grad, hr_r.grad), + ("orig_res", orig_f.grad, orig_r.grad), + ("h_post", hp_f.grad, hp_r.grad), + ("x", x_f.grad, x_r.grad), + ]: + torch.testing.assert_close( + gf, gr, atol=BWD_ATOL, rtol=BWD_RTOL, msg=f"backward mismatch on {name}" + ) + if with_bias: + torch.testing.assert_close( + bi_f.grad, bi_r.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on bias" + ) + + +class TestFusedHPostBDA: + """Public fused h_post_bda dispatch/fallback plus numerical correctness.""" + + def test_forward_uses_h_res_transpose(self): + from megatron.core.fusions.fused_mhc_kernels import fused_h_post_bda + + h_res, orig, h_post, x, expected = _h_post_bda_transpose_case() + out = fused_h_post_bda(h_res, orig, h_post, x, bias=None) + + torch.testing.assert_close(out, expected, atol=0.0, rtol=0.0) + + @pytest.mark.flaky_in_dev + @_require_cutile + @pytest.mark.parametrize("with_bias", [True, False]) + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256)]) + def test_fwd_bwd_vs_reference(self, s, b, n, C, with_bias): + """E2E: public fused fwd output and bwd grads must match the PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import fused_h_post_bda + + _info() + hr_data = _rand(s, b, n, n) + orig_data = _rand(s, b, n, C) + hp_data = _rand(s, b, n) + x_data = _rand(s, b, C) + bias_data = _rand(C) if with_bias else None + grad_out = _rand(s, b, n, C) + + def _make_inputs(): + hr = hr_data.clone().requires_grad_(True) + orig = orig_data.clone().requires_grad_(True) + hp = hp_data.clone().requires_grad_(True) + x = x_data.clone().requires_grad_(True) + bi = bias_data.clone().requires_grad_(True) if with_bias else None + return hr, orig, hp, x, bi + + # -- fused path -- + hr_f, orig_f, hp_f, x_f, bi_f = _make_inputs() + out_f = fused_h_post_bda(hr_f, orig_f, hp_f, x_f, bi_f) + out_f.backward(grad_out) + + # -- reference path -- + hr_r, orig_r, hp_r, x_r, bi_r = _make_inputs() + out_r = _ref_h_post_bda(hr_r, orig_r, hp_r, x_r, bi_r) + out_r.backward(grad_out) + + torch.testing.assert_close(out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + for name, gf, gr in [ + ("h_res", hr_f.grad, hr_r.grad), + ("orig_res", orig_f.grad, orig_r.grad), + ("h_post", hp_f.grad, hp_r.grad), + ("x", x_f.grad, x_r.grad), + ]: + torch.testing.assert_close( + gf, gr, atol=BWD_ATOL, rtol=BWD_RTOL, msg=f"backward mismatch on {name}" + ) + if with_bias: + torch.testing.assert_close( + bi_f.grad, bi_r.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on bias" + ) + + +class TestTritonHPostBDA: + """Tests for Triton h_post_bda kernels against PyTorch reference.""" + + @_require_triton + @pytest.mark.parametrize("with_bias", [True, False]) + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256), (64, 8, 4, 4096)]) + def test_fwd_vs_reference(self, s, b, n, C, with_bias): + """Triton hpb forward output must match the PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import _triton_h_post_bda_fwd + + _info() + hr_data = _rand(s, b, n, n) + orig_data = _rand(s, b, n, C) + hp_data = _rand(s, b, n) + x_data = _rand(s, b, C) + bias_data = _rand(C) if with_bias else None + + out_t = _triton_h_post_bda_fwd(hr_data, orig_data, hp_data, x_data, bias_data) + out_r = _ref_h_post_bda(hr_data, orig_data, hp_data, x_data, bias_data) + + torch.testing.assert_close(out_t, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + + @_require_triton + @pytest.mark.parametrize("with_bias", [True, False]) + @pytest.mark.parametrize( + "s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256), (64, 8, 4, 4096), (128, 1, 8, 7168)] + ) + def test_bwd_vs_reference(self, s, b, n, C, with_bias): + """Triton hpb backward grads must match the PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import _triton_h_post_bda_bwd + + _info() + hr_data = _rand(s, b, n, n) + orig_data = _rand(s, b, n, C) + hp_data = _rand(s, b, n) + x_data = _rand(s, b, C) + bias_data = _rand(C) if with_bias else None + grad_out = _rand(s, b, n, C) + + # -- Triton backward -- + g_hr_t, g_res_t, g_hp_t, g_x_t, g_bias_t = _triton_h_post_bda_bwd( + grad_out, hr_data, orig_data, hp_data, x_data, bias_data + ) + + # -- Reference backward via autograd -- + hr_r = hr_data.clone().requires_grad_(True) + orig_r = orig_data.clone().requires_grad_(True) + hp_r = hp_data.clone().requires_grad_(True) + x_r = x_data.clone().requires_grad_(True) + bi_r = bias_data.clone().requires_grad_(True) if with_bias else None + out_r = _ref_h_post_bda(hr_r, orig_r, hp_r, x_r, bi_r) + out_r.backward(grad_out) + + for name, gt, gr in [ + ("h_res", g_hr_t, hr_r.grad), + ("orig_res", g_res_t, orig_r.grad), + ("h_post", g_hp_t, hp_r.grad), + ("x", g_x_t, x_r.grad), + ]: + torch.testing.assert_close( + gt, gr, atol=BWD_ATOL, rtol=BWD_RTOL, msg=f"Triton backward mismatch on {name}" + ) + if with_bias: + torch.testing.assert_close( + g_bias_t, + bi_r.grad, + atol=BWD_ATOL, + rtol=BWD_RTOL, + msg="Triton backward mismatch on bias", + ) + + @_require_triton + @_require_cutile + @pytest.mark.parametrize("with_bias", [True, False]) + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256), (64, 8, 4, 4096)]) + def test_triton_vs_cutile(self, s, b, n, C, with_bias): + """Triton and cuTile backward must produce identical results.""" + from megatron.core.fusions.fused_mhc_kernels import ( + _cutile_h_post_bda_bwd, + _triton_h_post_bda_bwd, + ) + + _info() + hr_data = _rand(s, b, n, n) + orig_data = _rand(s, b, n, C) + hp_data = _rand(s, b, n) + x_data = _rand(s, b, C) + bias_data = _rand(C) if with_bias else None + grad_out = _rand(s, b, n, C) + + triton_out = _triton_h_post_bda_bwd( + grad_out, hr_data, orig_data, hp_data, x_data, bias_data + ) + cutile_out = _cutile_h_post_bda_bwd( + grad_out, hr_data, orig_data, hp_data, x_data, bias_data + ) + + for i, name in enumerate(["h_res", "orig_res", "h_post", "x", "bias"]): + if triton_out[i] is None: + continue + torch.testing.assert_close( + triton_out[i], + cutile_out[i], + atol=BWD_ATOL, + rtol=BWD_RTOL, + msg=f"Triton vs cuTile mismatch on {name}", + ) + + +class TestTritonHPostBDABwdE2EDebug: + """Debug: run E2E forward, then compare cuTile vs Triton backward per-output.""" + + @_require_triton + @_require_cutile + def test_e2e_inputs_no_nan(self): + """Feed actual E2E backward inputs to Triton kernel and check for NaN.""" + from megatron.core.fusions.fused_mhc_kernels import ( + _cutile_h_post_bda_bwd, + _triton_h_post_bda_bwd, + fused_h_aggregate, + fused_h_post_bda, + fused_proj_rms, + fused_sinkhorn, + ) + + s, b, n, C = 8, 4, 4, 1024 + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(n * n + 2 * n, n * C) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + # Run E2E forward to produce realistic backward inputs + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + x_2d = hs.reshape(s * b, n * C) + proj, r = fused_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post_val = h[..., n : 2 * n].sigmoid() * 2 + h_res = fused_sinkhorn(h[..., 2 * n :].view(s, b, n, n), sinkhorn_iters, eps) + _ = fused_h_aggregate(hs.view(s, b, n, C), h_pre) + output = fused_h_post_bda( + h_res, hs.view(s, b, n, C), h_post_val, layer_out_data, layer_bias_data + ) + go = torch.ones_like(output) + + # Capture inputs (detach from graph) + hr = h_res.detach() + orig = hs.view(s, b, n, C).detach() + hp = h_post_val.detach() + x = layer_out_data.detach() + bias = layer_bias_data.detach() + + # Compare cuTile vs Triton per-output + ct_out = _cutile_h_post_bda_bwd(go, hr, orig, hp, x, bias) + tr_out = _triton_h_post_bda_bwd(go, hr, orig, hp, x, bias) + + names = ["g_hr", "g_res", "g_hp", "g_x", "g_bias"] + for name, ct_t, tr_t in zip(names, ct_out, tr_out): + if tr_t is None: + continue + assert not tr_t.isnan().any(), f"Triton {name} has NaN" + assert not tr_t.isinf().any(), f"Triton {name} has Inf" + torch.testing.assert_close( + tr_t, + ct_t, + atol=BWD_ATOL, + rtol=BWD_RTOL, + msg=f"Triton vs cuTile mismatch on {name} (E2E inputs)", + ) + + +# ============================================================================ +# Triton: Sinkhorn +# ============================================================================ + + +class TestTritonSinkhorn: + @_require_triton + @pytest.mark.parametrize("s,b,n,iters", [(2, 4, 4, 5), (1, 1, 2, 10), (8, 4, 4, 20)]) + def test_fwd_bwd_vs_reference(self, s, b, n, iters): + from megatron.core.fusions.fused_mhc_kernels import triton_fused_sinkhorn + + eps = 1e-6 + data = _rand(s, b, n, n) + grad_out = _rand(s, b, n, n) + + inp_f = data.clone().requires_grad_(True) + out_f = triton_fused_sinkhorn(inp_f, iters, eps) + out_f.backward(grad_out) + + inp_r = data.clone().requires_grad_(True) + out_r = _ref_sinkhorn(inp_r, iters, eps) + out_r.backward(grad_out) + + torch.testing.assert_close(out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(inp_f.grad, inp_r.grad, atol=BWD_ATOL, rtol=BWD_RTOL) + + @_require_triton + @_require_cutile + @pytest.mark.parametrize("s,b,n,iters", [(2, 4, 4, 5)]) + def test_triton_vs_cutile(self, s, b, n, iters): + from megatron.core.fusions.fused_mhc_kernels import ( + _cutile_sinkhorn_bwd, + _cutile_sinkhorn_fwd, + triton_fused_sinkhorn, + ) + + eps = 1e-6 + data = _rand(s, b, n, n) + grad_out = _rand(s, b, n, n) + + inp_t = data.clone().requires_grad_(True) + out_t = triton_fused_sinkhorn(inp_t, iters, eps) + out_t.backward(grad_out) + + out_c, M_init = _cutile_sinkhorn_fwd(data.clone(), iters, eps) + grad_c = _cutile_sinkhorn_bwd(grad_out, M_init, iters, eps) + + torch.testing.assert_close(out_t, out_c, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(inp_t.grad, grad_c, atol=BWD_ATOL, rtol=BWD_RTOL) + + +# ============================================================================ +# Proj RMS +# ============================================================================ + + +class TestNativeProjRms: + """Tests for native_proj_rms.""" + + @pytest.mark.parametrize("M,N,K", [(256, 20, 4096), (64, 8, 512)]) + def test_fwd_bwd_vs_torch_reference(self, M, N, K): + _info() + eps = 1e-6 + x_data = _rand(M, K) + w_data = _rand(N, K) + grad_proj = _rand(M, N) + grad_r = _rand(M, 1) + + xf = x_data.clone().requires_grad_(True) + wf = w_data.clone().requires_grad_(True) + proj_f, r_f = native_proj_rms(xf, wf, eps) + (proj_f * grad_proj + r_f * grad_r).sum().backward() + + xr = x_data.clone().requires_grad_(True) + wr = w_data.clone().requires_grad_(True) + proj_r, r_r = _ref_proj_rms(xr, wr, eps) + (proj_r * grad_proj + r_r * grad_r).sum().backward() + + torch.testing.assert_close(proj_f, proj_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(r_f, r_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close( + xf.grad, xr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on x" + ) + torch.testing.assert_close( + wf.grad, wr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on weight" + ) + + +class TestFusedProjRms: + """Public fused proj_rms dispatch/fallback plus numerical correctness.""" + + @pytest.mark.flaky_in_dev + @_require_cutile + @pytest.mark.parametrize("M,N,K", [(256, 20, 4096), (64, 8, 512)]) + def test_fwd_bwd_vs_reference(self, M, N, K): + """E2E: public fused fwd output and bwd grads must match the PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import fused_proj_rms + + _info() + eps = 1e-6 + x_data = _rand(M, K) + w_data = _rand(N, K) + grad_proj = _rand(M, N) + grad_r = _rand(M, 1) + + # -- fused path -- + xf = x_data.clone().requires_grad_(True) + wf = w_data.clone().requires_grad_(True) + proj_f, r_f = fused_proj_rms(xf, wf, eps) + (proj_f * grad_proj + r_f * grad_r).sum().backward() + + # -- reference path -- + xr = x_data.clone().requires_grad_(True) + wr = w_data.clone().requires_grad_(True) + proj_r, r_r = _ref_proj_rms(xr, wr, eps) + (proj_r * grad_proj + r_r * grad_r).sum().backward() + + torch.testing.assert_close(proj_f, proj_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(r_f, r_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close( + xf.grad, xr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on x" + ) + torch.testing.assert_close( + wf.grad, wr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on weight" + ) + + +# ============================================================================ +# Proj RMS + Compute H (fused) +# ============================================================================ + + +class TestFusedProjRmsComputeH: + """Public fused proj_rms_compute_h dispatch/fallback plus numerical correctness.""" + + @pytest.mark.parametrize("M,n,K", [(256, 4, 4096), (64, 2, 512), (128, 4, 2048)]) + def test_fwd_bwd_vs_reference(self, M, n, K): + """E2E: public fused fwd output and bwd grads must match the PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import fused_proj_rms_compute_h + + _info() + N = n * n + 2 * n + eps = 1e-6 + x_data = _rand(M, K) + w_data = _rand(N, K) + ap_data = _rand(1) + apo_data = _rand(1) + ar_data = _rand(1) + bias_data = _rand(N) + grad_y = _rand(M, N) + grad_h_pre = grad_y[:, :n] + grad_h_post = grad_y[:, n : 2 * n] + grad_h_res = grad_y[:, 2 * n :] + grad_r = _rand(M, 1) + + def _make_inputs(): + return ( + x_data.clone().requires_grad_(True), + w_data.clone().requires_grad_(True), + ap_data.clone().requires_grad_(True), + apo_data.clone().requires_grad_(True), + ar_data.clone().requires_grad_(True), + bias_data.clone().requires_grad_(True), + ) + + # -- fused path -- + xf, wf, apf, apof, arf, bf = _make_inputs() + h_pre_f, h_post_f, h_res_f, r_f = fused_proj_rms_compute_h( + xf, wf, apf, apof, arf, bf, n, eps + ) + loss_f = ( + (h_pre_f * grad_h_pre).sum() + + (h_post_f * grad_h_post).sum() + + (h_res_f * grad_h_res).sum() + + (r_f * grad_r).sum() + ) + loss_f.backward() + + # -- reference path -- + xr, wr, apr, apor, arr, br = _make_inputs() + h_pre_r, h_post_r, h_res_r, r_r = _ref_proj_rms_compute_h( + xr, wr, apr, apor, arr, br, n, eps + ) + loss_r = ( + (h_pre_r * grad_h_pre).sum() + + (h_post_r * grad_h_post).sum() + + (h_res_r * grad_h_res).sum() + + (r_r * grad_r).sum() + ) + loss_r.backward() + + torch.testing.assert_close( + h_pre_f, h_pre_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_pre mismatch" + ) + torch.testing.assert_close( + h_post_f, h_post_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post mismatch" + ) + torch.testing.assert_close( + h_res_f, h_res_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_res mismatch" + ) + torch.testing.assert_close(r_f, r_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="r mismatch") + torch.testing.assert_close( + xf.grad, xr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on x" + ) + torch.testing.assert_close( + wf.grad, wr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on weight" + ) + _assert_cosine_similar( + apf.grad, apr.grad, COSINE_SIM_THRESH, msg="backward mismatch on alpha_pre" + ) + _assert_cosine_similar( + apof.grad, apor.grad, COSINE_SIM_THRESH, msg="backward mismatch on alpha_post" + ) + _assert_cosine_similar( + arf.grad, arr.grad, COSINE_SIM_THRESH, msg="backward mismatch on alpha_res" + ) + _assert_cosine_similar(bf.grad, br.grad, COSINE_SIM_THRESH, msg="backward mismatch on bias") + + +# ============================================================================ +# End-to-end pipeline (all four kernels chained) +# ============================================================================ + + +class TestEndToEndNative: + """Full mHC pipeline using native modules. + + proj_rms -> compute_h -> sinkhorn -> aggregate -> h_post_bda. + Compares the native modules against inline PyTorch reference. + """ + + def test_full_pipeline_fwd_bwd(self): + _info() + s, b, n, C = 2, 4, 4, 1024 + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(n * n + 2 * n, n * C) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + def _run_native_modules(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + proj, r = native_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = native_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = native_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = native_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + def _run_inline_ref(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + proj, r = _ref_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = _ref_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = _ref_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = _ref_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + out_m, agg_m, grad_m = _run_native_modules() + out_r, agg_r, grad_r = _run_inline_ref() + + torch.testing.assert_close( + agg_m, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="aggregated output mismatch" + ) + torch.testing.assert_close( + out_m, out_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post_bda output mismatch" + ) + _assert_cosine_similar( + grad_m, grad_r, COSINE_SIM_THRESH, msg="hidden_states grad (E2E backward)" + ) + + +class TestEndToEndFused: + """Full mHC pipeline using the public fused API.""" + + @_require_cutile + def test_full_pipeline_fwd_bwd(self): + from megatron.core.fusions.fused_mhc_kernels import ( + fused_h_aggregate, + fused_h_post_bda, + fused_proj_rms, + fused_sinkhorn, + ) + + _info() + s, b, n, C = 8, 4, 4, 1024 + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(n * n + 2 * n, n * C) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + def _run_fused(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + proj, r = fused_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = fused_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = fused_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = fused_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return proj.detach(), r.detach(), output.detach(), aggregated.detach(), hs.grad.clone() + + def _run_ref(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + proj, r = _ref_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = _ref_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = _ref_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = _ref_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return proj.detach(), r.detach(), output.detach(), aggregated.detach(), hs.grad.clone() + + proj_f, r_f, out_f, agg_f, grad_f = _run_fused() + proj_r, r_r, out_r, agg_r, grad_r = _run_ref() + + torch.testing.assert_close(proj_f, proj_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(r_f, r_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(agg_f, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close( + out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post_bda output mismatch" + ) + _assert_cosine_similar( + grad_f, grad_r, COSINE_SIM_THRESH, msg="hidden_states grad (E2E backward)" + ) + + def test_full_pipeline_fused_compute_h(self): + """E2E: fused proj_rms_compute_h replaces separate proj_rms + compute_h.""" + from megatron.core.fusions.fused_mhc_kernels import ( + fused_h_aggregate, + fused_h_post_bda, + fused_proj_rms_compute_h, + fused_sinkhorn, + ) + + _info() + s, b, n, C = 8, 4, 4, 1024 + N = n * n + 2 * n + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(N, n * C) + ap_data = _rand(1) + apo_data = _rand(1) + ar_data = _rand(1) + bias_data = _rand(N) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + def _run_fused_compute_h(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + ap = ap_data.clone().requires_grad_(True) + apo = apo_data.clone().requires_grad_(True) + ar = ar_data.clone().requires_grad_(True) + bias_p = bias_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + h_pre, h_post, h_res_logits, _ = fused_proj_rms_compute_h( + x_2d, w, ap, apo, ar, bias_p, n, eps + ) + + h_pre = h_pre.view(s, b, n) + h_post = h_post.view(s, b, n) + h_res_logits = h_res_logits.view(s, b, n, n) + h_res = fused_sinkhorn(h_res_logits, sinkhorn_iters, eps) + + aggregated = fused_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = fused_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + def _run_ref(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + ap = ap_data.clone().requires_grad_(True) + apo = apo_data.clone().requires_grad_(True) + ar = ar_data.clone().requires_grad_(True) + bias_p = bias_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + h_pre, h_post, h_res_logits, _ = _ref_proj_rms_compute_h( + x_2d, w, ap, apo, ar, bias_p, n, eps + ) + + h_pre = h_pre.view(s, b, n) + h_post = h_post.view(s, b, n) + h_res_logits = h_res_logits.view(s, b, n, n) + h_res = _ref_sinkhorn(h_res_logits, sinkhorn_iters, eps) + + aggregated = _ref_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = _ref_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + out_f, agg_f, grad_f = _run_fused_compute_h() + out_r, agg_r, grad_r = _run_ref() + + torch.testing.assert_close( + agg_f, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="aggregated mismatch" + ) + torch.testing.assert_close( + out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post_bda output mismatch" + ) + _assert_cosine_similar( + grad_f, + grad_r, + COSINE_SIM_THRESH, + msg="hidden_states grad (E2E backward, fused compute_h)", + ) + + +# ============================================================================ +# fused_add_3 kernel tests +# ============================================================================ + + +class TestFusedAdd3: + """Tests for fused_add_3 (torch.compile backend, no cuTile dependency).""" + + def test_fused_add_3_forward_bf16(self): + """fused_add_3 matches a + b + c for bf16 tensors.""" + from megatron.core.fusions.fused_mhc_kernels import fused_add_3 + + a = torch.randn(128, 256, dtype=DTYPE, device=DEVICE) + b = torch.randn(128, 256, dtype=DTYPE, device=DEVICE) + c = torch.randn(128, 256, dtype=DTYPE, device=DEVICE) + result = fused_add_3(a, b, c) + expected = (a.float() + b.float() + c.float()).to(DTYPE) + torch.testing.assert_close(result, expected, atol=FWD_ATOL, rtol=FWD_RTOL) + + def test_fused_add_3_forward_fp32(self): + """fused_add_3 matches a + b + c for fp32 tensors.""" + from megatron.core.fusions.fused_mhc_kernels import fused_add_3 + + a = torch.randn(128, 256, dtype=torch.float32, device=DEVICE) + b = torch.randn(128, 256, dtype=torch.float32, device=DEVICE) + c = torch.randn(128, 256, dtype=torch.float32, device=DEVICE) + result = fused_add_3(a, b, c) + expected = a + b + c + torch.testing.assert_close(result, expected, atol=1e-5, rtol=1e-5) + + def test_fused_add_3_large_tensor(self): + """fused_add_3 handles large tensors.""" + from megatron.core.fusions.fused_mhc_kernels import fused_add_3 + + a = torch.randn(8192, 4096, dtype=DTYPE, device=DEVICE) + b = torch.randn(8192, 4096, dtype=DTYPE, device=DEVICE) + c = torch.randn(8192, 4096, dtype=DTYPE, device=DEVICE) + result = fused_add_3(a, b, c) + expected = (a.float() + b.float() + c.float()).to(DTYPE) + torch.testing.assert_close(result, expected, atol=FWD_ATOL, rtol=FWD_RTOL) + + def test_fused_add_3_gradient(self): + """fused_add_3 produces correct gradients.""" + from megatron.core.fusions.fused_mhc_kernels import fused_add_3 + + a = torch.randn(64, 128, dtype=torch.float32, device=DEVICE, requires_grad=True) + b = torch.randn(64, 128, dtype=torch.float32, device=DEVICE, requires_grad=True) + c = torch.randn(64, 128, dtype=torch.float32, device=DEVICE, requires_grad=True) + result = fused_add_3(a, b, c) + result.sum().backward() + torch.testing.assert_close(a.grad, torch.ones_like(a)) + torch.testing.assert_close(b.grad, torch.ones_like(b)) + torch.testing.assert_close(c.grad, torch.ones_like(c)) + + +# ============================================================================ +# End-to-end pipeline with BroadcastTensorFused +# ============================================================================ + + +class TestEndToEndNativeBroadcast: + """Full mHC pipeline using native modules + BroadcastTensorFused. + + Same pipeline as TestEndToEndNative but hidden_states is split via + BroadcastTensorFused so each consumer (proj_rms/compute_h, aggregate, + h_post_bda) gets a distinct autograd graph node. Verifies gradient + correctness versus the inline reference that uses the tensor directly. + """ + + def test_full_pipeline_fwd_bwd(self): + from megatron.core.transformer.hyper_connection import ( + BroadcastTensorFused, + native_fused_add_3, + ) + + _info() + s, b, n, C = 2, 4, 4, 1024 + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(n * n + 2 * n, n * C) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + def _run_broadcast(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + # Split via BroadcastTensorFused + hs_map, hs_agg, hs_res = BroadcastTensorFused.apply(hs, native_fused_add_3) + + x_2d = hs_map.reshape(s * b, n * C) + proj, r = native_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = native_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = native_h_aggregate(hs_agg.view(s, b, n, C), h_pre) + + output = native_h_post_bda( + h_res, hs_res.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + def _run_inline_ref(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + proj, r = _ref_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = _ref_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = _ref_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = _ref_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + out_b, agg_b, grad_b = _run_broadcast() + out_r, agg_r, grad_r = _run_inline_ref() + + torch.testing.assert_close( + agg_b, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="aggregated mismatch (broadcast)" + ) + torch.testing.assert_close( + out_b, out_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post_bda output mismatch (broadcast)" + ) + _assert_cosine_similar( + grad_b, grad_r, COSINE_SIM_THRESH, msg="hidden_states grad (E2E backward, broadcast)" + ) + + +class TestEndToEndFusedBroadcast: + """Full mHC pipeline through public fused dispatch + BroadcastTensorFused.""" + + def test_full_pipeline_fwd_bwd(self): + from megatron.core.fusions.fused_mhc_kernels import ( + fused_add_3, + fused_h_aggregate, + fused_h_post_bda, + fused_proj_rms, + fused_sinkhorn, + ) + from megatron.core.transformer.hyper_connection import BroadcastTensorFused + + _info() + s, b, n, C = 8, 4, 4, 1024 + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(n * n + 2 * n, n * C) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + def _run_fused_broadcast(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + hs_map, hs_agg, hs_res = BroadcastTensorFused.apply(hs, fused_add_3) + + x_2d = hs_map.reshape(s * b, n * C) + proj, r = fused_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = fused_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = fused_h_aggregate(hs_agg.view(s, b, n, C), h_pre) + + output = fused_h_post_bda( + h_res, hs_res.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + def _run_ref(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + proj, r = _ref_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = _ref_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = _ref_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = _ref_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + out_f, agg_f, grad_f = _run_fused_broadcast() + out_r, agg_r, grad_r = _run_ref() + + torch.testing.assert_close( + agg_f, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="aggregated mismatch (fused broadcast)" + ) + torch.testing.assert_close( + out_f, + out_r, + atol=FWD_ATOL, + rtol=FWD_RTOL, + msg="h_post_bda output mismatch (fused broadcast)", + ) + _assert_cosine_similar( + grad_f, + grad_r, + COSINE_SIM_THRESH, + msg="hidden_states grad (E2E backward, fused broadcast)", + ) + + def test_full_pipeline_fused_compute_h_broadcast(self): + """E2E: fused proj_rms_compute_h + BroadcastTensorFused.""" + from megatron.core.fusions.fused_mhc_kernels import ( + fused_add_3, + fused_h_aggregate, + fused_h_post_bda, + fused_proj_rms_compute_h, + fused_sinkhorn, + ) + from megatron.core.transformer.hyper_connection import BroadcastTensorFused + + _info() + s, b, n, C = 8, 4, 4, 1024 + N = n * n + 2 * n + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(N, n * C) + ap_data = _rand(1) + apo_data = _rand(1) + ar_data = _rand(1) + bias_data = _rand(N) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + def _run_fused_compute_h_broadcast(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + ap = ap_data.clone().requires_grad_(True) + apo = apo_data.clone().requires_grad_(True) + ar = ar_data.clone().requires_grad_(True) + bias_p = bias_data.clone().requires_grad_(True) + + hs_map, hs_agg, hs_res = BroadcastTensorFused.apply(hs, fused_add_3) + + x_2d = hs_map.reshape(s * b, n * C) + h_pre, h_post, h_res_logits, _ = fused_proj_rms_compute_h( + x_2d, w, ap, apo, ar, bias_p, n, eps + ) + + h_pre = h_pre.view(s, b, n) + h_post = h_post.view(s, b, n) + h_res_logits = h_res_logits.view(s, b, n, n) + h_res = fused_sinkhorn(h_res_logits, sinkhorn_iters, eps) + + aggregated = fused_h_aggregate(hs_agg.view(s, b, n, C), h_pre) + + output = fused_h_post_bda( + h_res, hs_res.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + def _run_ref(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + ap = ap_data.clone().requires_grad_(True) + apo = apo_data.clone().requires_grad_(True) + ar = ar_data.clone().requires_grad_(True) + bias_p = bias_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + h_pre, h_post, h_res_logits, _ = _ref_proj_rms_compute_h( + x_2d, w, ap, apo, ar, bias_p, n, eps + ) + + h_pre = h_pre.view(s, b, n) + h_post = h_post.view(s, b, n) + h_res_logits = h_res_logits.view(s, b, n, n) + h_res = _ref_sinkhorn(h_res_logits, sinkhorn_iters, eps) + + aggregated = _ref_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = _ref_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + out_f, agg_f, grad_f = _run_fused_compute_h_broadcast() + out_r, agg_r, grad_r = _run_ref() + + torch.testing.assert_close( + agg_f, + agg_r, + atol=FWD_ATOL, + rtol=FWD_RTOL, + msg="aggregated mismatch (fused compute_h broadcast)", + ) + torch.testing.assert_close( + out_f, + out_r, + atol=FWD_ATOL, + rtol=FWD_RTOL, + msg="h_post_bda output mismatch (fused compute_h broadcast)", + ) + _assert_cosine_similar( + grad_f, + grad_r, + COSINE_SIM_THRESH, + msg="hidden_states grad (E2E backward, fused compute_h broadcast)", + ) diff --git a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py index 810f48092ee..04f09c82b45 100644 --- a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py +++ b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py @@ -16,12 +16,16 @@ try: from megatron.core.fusions.fused_mla_yarn_rope_apply import ( - fused_apply_mla_rope_for_kv, fused_apply_mla_rope_for_q, + fused_mla_rope_inplace, + fused_mla_rope_kv_split, + fused_mla_rope_out_of_place, ) -except: - fused_apply_mla_rope_for_kv = None +except Exception: fused_apply_mla_rope_for_q = None + fused_mla_rope_inplace = None + fused_mla_rope_kv_split = None + fused_mla_rope_out_of_place = None def dtype_tols(dtype): @@ -54,7 +58,9 @@ def test_packed_freqs_returns_offset_mapped_output_for_context_parallel(self): t = torch.randn(4, 2, 8) freqs = torch.randn(8, 1, 1, 8) - out = rope_utils_module._apply_rotary_pos_emb_thd(t, cu_seqlens, freqs, cp_group=cp_group) + out = rope_utils_module._apply_rotary_pos_emb_thd( + t, cu_seqlens, freqs, cp_group=cp_group, max_seqlen=4 + ) expected_freqs = torch.cat([freqs[0:1], freqs[3:4], freqs[4:5], freqs[7:8]], dim=0) expected = rope_utils_module._apply_rotary_pos_emb_bshd( @@ -69,7 +75,9 @@ def test_max_seqlen_freqs_returns_sequence_mapped_output_for_context_parallel(se t = torch.randn(4, 2, 8) freqs = torch.randn(4, 1, 1, 8) - out = rope_utils_module._apply_rotary_pos_emb_thd(t, cu_seqlens, freqs, cp_group=cp_group) + out = rope_utils_module._apply_rotary_pos_emb_thd( + t, cu_seqlens, freqs, cp_group=cp_group, max_seqlen=4 + ) expected_freqs = torch.cat([freqs[1:2], freqs[2:3]], dim=0) expected_slices = [] @@ -83,9 +91,39 @@ def test_max_seqlen_freqs_returns_sequence_mapped_output_for_context_parallel(se torch.testing.assert_close(out, expected) + def test_missing_max_seqlen_preserves_legacy_packed_freq_mapping(self): + cp_group = FakeCPGroup(size=2, rank=0) + cu_seqlens = torch.tensor([0, 4, 8], dtype=torch.int32) + t = torch.randn(4, 2, 8) + freqs = torch.randn(8, 1, 1, 8) -def _test_fused_apply_mla_rope_for_q(input_format): - assert fused_apply_mla_rope_for_q is not None + legacy_out = rope_utils_module._apply_rotary_pos_emb_thd( + t, cu_seqlens, freqs, cp_group=cp_group + ) + explicit_out = rope_utils_module._apply_rotary_pos_emb_thd( + t, cu_seqlens, freqs, cp_group=cp_group, max_seqlen=4 + ) + + torch.testing.assert_close(legacy_out, explicit_out) + + +class _SaveOutputForBackward(torch.autograd.Function): + """Minimal stand-in for a kernel whose backward consumes its output.""" + + @staticmethod + def forward(ctx, tensor): + output = tensor.clone() + ctx.save_for_backward(output) + return output + + @staticmethod + def backward(ctx, _grad_output): + (saved_output,) = ctx.saved_tensors + return saved_output + + +def _test_fused_mla_rope_inplace(input_format, inverse=False, remove_interleaving=False): + assert fused_mla_rope_inplace is not None num_heads = 32 q_dim = 128 emb_dim = 64 @@ -97,6 +135,7 @@ def _test_fused_apply_mla_rope_for_q(input_format): multi_latent_attention=True, ) + max_seqlen = None if input_format == "sbhd": cu_seqlens = None seqlen = 1024 @@ -142,15 +181,25 @@ def _test_fused_apply_mla_rope_for_q(input_format): freqs, transformer_config, cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, mscale=mscale, cp_group=FakeCPGroup(), mla_rotary_interleaved=True, + inverse=inverse, + mla_output_remove_interleaving=remove_interleaving, ) pytorch_output = torch.concat([no_pe, pe_output], dim=-1) pytorch_output.backward(pytorch_bwd_input, retain_graph=True) - fused_output = fused_apply_mla_rope_for_q( - fused_fwd_input, cos, sin, q_dim, emb_dim, cu_seqlens_q=cu_seqlens + fused_output = fused_mla_rope_inplace( + fused_fwd_input, + cos, + sin, + q_dim, + emb_dim, + cu_seqlens_q=cu_seqlens, + inverse=inverse, + remove_interleaving=remove_interleaving, ) fused_output.backward(fused_bwd_input, retain_graph=True) @@ -169,8 +218,8 @@ def _test_fused_apply_mla_rope_for_q(input_format): ) -def _test_fused_apply_mla_rope_for_kv(input_format): - assert fused_apply_mla_rope_for_kv is not None +def _test_fused_mla_rope_kv_split(input_format, remove_interleaving=False): + assert fused_mla_rope_kv_split is not None num_heads = 32 k_dim = 128 v_dim = 128 @@ -183,6 +232,7 @@ def _test_fused_apply_mla_rope_for_kv(input_format): multi_latent_attention=True, ) + max_seqlen = None if input_format == "sbhd": cu_seqlens = None seqlen = 1024 @@ -241,9 +291,11 @@ def _test_fused_apply_mla_rope_for_kv(input_format): freqs, transformer_config, cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, mscale=mscale, cp_group=FakeCPGroup(), mla_rotary_interleaved=True, + mla_output_remove_interleaving=remove_interleaving, ) if input_format == "sbhd": pe_output = pe_output.expand(-1, -1, num_heads, -1) @@ -255,7 +307,7 @@ def _test_fused_apply_mla_rope_for_kv(input_format): (pytorch_k_output, pytorch_v_output), (pytorch_bwd_k_input, pytorch_bwd_v_input) ) - fused_k_output, fused_v_output = fused_apply_mla_rope_for_kv( + fused_k_output, fused_v_output = fused_mla_rope_kv_split( fused_fwd_kv_input, fused_fwd_emb_input, cos, @@ -264,6 +316,7 @@ def _test_fused_apply_mla_rope_for_kv(input_format): k_dim, v_dim, cu_seqlens_kv=cu_seqlens, + remove_interleaving=remove_interleaving, ) torch.autograd.backward( (fused_k_output, fused_v_output), (fused_bwd_k_input, fused_bwd_v_input) @@ -301,13 +354,136 @@ def _test_fused_apply_mla_rope_for_kv(input_format): @pytest.mark.skipif(not is_torch_min_version("2.5.0"), reason="Requires PyTorch >= 2.5.0") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.parametrize("input_format", ["sbhd", "thd"]) -class TestFusedApplyMLARope: +class TestFusedMLARope: @pytest.mark.flaky_in_dev - def test_forward_backward_for_q(self, input_format): - _test_fused_apply_mla_rope_for_q(input_format) + @pytest.mark.parametrize("inverse", [False, True]) + @pytest.mark.parametrize("remove_interleaving", [False, True]) + def test_inplace_forward_backward(self, input_format, inverse, remove_interleaving): + _test_fused_mla_rope_inplace( + input_format, inverse=inverse, remove_interleaving=remove_interleaving + ) + + @pytest.mark.parametrize("remove_interleaving", [False, True]) + def test_kv_split_forward_backward(self, input_format, remove_interleaving): + _test_fused_mla_rope_kv_split(input_format, remove_interleaving=remove_interleaving) + + +@pytest.mark.experimental +@pytest.mark.internal +@pytest.mark.skipif(not is_torch_min_version("2.5.0"), reason="Requires PyTorch >= 2.5.0") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize("input_format", ["sbhd", "thd"]) +def test_out_of_place_inverse_rope_preserves_upstream_saved_output(input_format): + """Post-attention inverse RoPE must not overwrite an output saved for backward.""" + assert fused_mla_rope_out_of_place is not None + seqlen = 32 + batch_size = 1 + num_heads = 2 + nope_dim = 16 + emb_dim = 64 + dtype = torch.bfloat16 + + yarn_rope = YarnRotaryEmbedding(emb_dim, original_max_position_embeddings=seqlen) + freqs, mscale = yarn_rope(seqlen, 0) + cos = (torch.cos(freqs) * mscale).to(dtype) + sin = (torch.sin(freqs) * mscale).to(dtype) + + if input_format == "sbhd": + shape = (seqlen, batch_size, num_heads, nope_dim + emb_dim) + cu_seqlens = None + else: + shape = (2 * seqlen, num_heads, nope_dim + emb_dim) + cu_seqlens = torch.tensor([0, seqlen, 2 * seqlen], dtype=torch.int32, device="cuda") + + unsafe_source = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True) + unsafe_attention_output = _SaveOutputForBackward.apply(unsafe_source) + unsafe_reference = unsafe_attention_output.detach().clone() + unsafe_inverse_output = fused_mla_rope_inplace( + unsafe_attention_output, + cos, + sin, + nope_dim, + emb_dim, + cu_seqlens_q=cu_seqlens, + inverse=True, + remove_interleaving=True, + ) + + assert unsafe_inverse_output.data_ptr() == unsafe_attention_output.data_ptr() + assert not torch.equal(unsafe_attention_output, unsafe_reference) + + source = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True) + attention_output = _SaveOutputForBackward.apply(source) + saved_reference = attention_output.detach().clone() + + inverse_output = fused_mla_rope_out_of_place( + attention_output, + cos, + sin, + nope_dim, + emb_dim, + cu_seqlens_q=cu_seqlens, + inverse=True, + remove_interleaving=True, + ) + expected_inverse_output = fused_mla_rope_inplace( + saved_reference.clone(), + cos, + sin, + nope_dim, + emb_dim, + cu_seqlens_q=cu_seqlens, + inverse=True, + remove_interleaving=True, + ) + + assert inverse_output.data_ptr() != attention_output.data_ptr() + torch.testing.assert_close(attention_output, saved_reference, rtol=0, atol=0) + torch.testing.assert_close(inverse_output, expected_inverse_output, rtol=0, atol=0) + + inverse_output.backward(torch.randn_like(inverse_output).contiguous()) + torch.testing.assert_close(source.grad, saved_reference, rtol=0, atol=0) + + +@pytest.mark.experimental +@pytest.mark.internal +@pytest.mark.skipif(not is_torch_min_version("2.5.0"), reason="Requires PyTorch >= 2.5.0") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize("input_format", ["sbhd", "thd"]) +def test_legacy_query_api_remains_in_place(input_format): + """The legacy API keeps its original mutation behavior and allocation profile.""" + assert fused_apply_mla_rope_for_q is not None + seqlen = 32 + batch_size = 1 + num_heads = 2 + nope_dim = 16 + emb_dim = 64 + dtype = torch.bfloat16 + + yarn_rope = YarnRotaryEmbedding(emb_dim, original_max_position_embeddings=seqlen) + freqs, mscale = yarn_rope(seqlen, 0) + cos = (torch.cos(freqs) * mscale).to(dtype) + sin = (torch.sin(freqs) * mscale).to(dtype) + + if input_format == "sbhd": + shape = (seqlen, batch_size, num_heads, nope_dim + emb_dim) + cu_seqlens = None + else: + shape = (2 * seqlen, num_heads, nope_dim + emb_dim) + cu_seqlens = torch.tensor([0, seqlen, 2 * seqlen], dtype=torch.int32, device="cuda") + + query = torch.randn(shape, dtype=dtype, device="cuda") + reference = query.clone() + expected = fused_mla_rope_inplace( + reference.clone(), cos, sin, nope_dim, emb_dim, cu_seqlens_q=cu_seqlens + ) + output = fused_apply_mla_rope_for_q( + query, cos, sin, qk_head_dim=nope_dim, emb_dim=emb_dim, cu_seqlens_q=cu_seqlens + ) - def test_forward_backward_for_kv(self, input_format): - _test_fused_apply_mla_rope_for_kv(input_format) + assert output.data_ptr() == query.data_ptr() + assert not torch.equal(query, reference) + torch.testing.assert_close(output, expected, rtol=0, atol=0) class TestApplyRotaryPosEmbMlaFusionConflict: diff --git a/tests/unit_tests/fusions/test_swiglu_fusion.py b/tests/unit_tests/fusions/test_swiglu_fusion.py index c72679cd047..1ba9b3eb891 100644 --- a/tests/unit_tests/fusions/test_swiglu_fusion.py +++ b/tests/unit_tests/fusions/test_swiglu_fusion.py @@ -1,7 +1,51 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + import pytest import torch +import torch.nn.functional as F from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl, weighted_bias_swiglu_impl +from megatron.core.transformer.transformer_config import TransformerConfig + + +def _clamped_swiglu_config(**kwargs): + defaults = dict( + num_layers=1, + hidden_size=16, + num_attention_heads=4, + num_moe_experts=4, + gated_linear_unit=True, + activation_func=F.silu, + activation_func_clamp_value=10.0, + ) + return TransformerConfig(**(defaults | kwargs)) + + +def test_clamped_swiglu_config_accepts_positive_moe_clamp(): + assert _clamped_swiglu_config().activation_func_clamp_value == 10.0 + + +@pytest.mark.parametrize("clamp_value", [0.0, -1.0, float("nan"), float("inf"), float("-inf")]) +def test_clamped_swiglu_config_requires_positive_clamp(clamp_value): + with pytest.raises(ValueError, match="greater than zero"): + _clamped_swiglu_config(activation_func_clamp_value=clamp_value) + + +def test_clamped_swiglu_config_rejects_linear_offset(): + with pytest.raises(ValueError, match="glu_linear_offset must be zero"): + _clamped_swiglu_config(glu_linear_offset=1.0) + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"num_moe_experts": None}, "only supported with MoE"), + ({"use_te_activation_func": True}, "use_te_activation_func must be False"), + ], +) +def test_clamped_swiglu_config_rejects_unsupported_paths(kwargs, match): + with pytest.raises(ValueError, match=match): + _clamped_swiglu_config(**kwargs) @pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float32]) @@ -39,3 +83,120 @@ def test_weighted_bias_swiglu(input_dtype): assert weights_2.grad.dtype == weights.grad.dtype if input_dtype == torch.float32: assert torch.allclose(weights.grad, weights_2.grad, **tols) + + +@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float32]) +def test_clamped_weighted_bias_swiglu(input_dtype): + clamp_value = 10.0 + + if input_dtype == torch.float32: + tols = dict(rtol=1.0e-6, atol=1.0e-6) + elif input_dtype == torch.bfloat16: + tols = dict(rtol=2.0e-2, atol=1.0e-3) + else: + raise ValueError(f"Invalid input dtype: {input_dtype}") + + x = (torch.randn(16, 64, dtype=input_dtype, device="cuda") * 5.0).requires_grad_(True) + weights = torch.randn(16, 1, dtype=torch.float32, device="cuda", requires_grad=True) + bwd_input = torch.randn(16, 32, dtype=input_dtype, device="cuda") + + # Reference: clamp and activate in FP32, then restore the input dtype. + y_1, y_2 = torch.chunk(x.to(torch.float32), 2, -1) + y = ( + F.silu(y_1.clamp(min=None, max=clamp_value)) + * y_2.clamp(min=-clamp_value, max=clamp_value) + * weights + ).to(input_dtype) + y.backward(bwd_input) + + x_fused = x.detach().clone().requires_grad_(True) + weights_fused = weights.detach().clone().requires_grad_(True) + y_fused = weighted_bias_swiglu_impl(x_fused, None, weights_fused, clamp_value=clamp_value) + y_fused.backward(bwd_input.detach().clone()) + + assert y_fused.dtype == y.dtype + assert torch.allclose(y, y_fused, **tols) + assert x_fused.grad.dtype == x.grad.dtype + assert torch.allclose(x.grad, x_fused.grad, **tols) + assert weights_fused.grad.dtype == weights.grad.dtype + if input_dtype == torch.float32: + assert torch.allclose(weights.grad, weights_fused.grad, **tols) + + +@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_clamped_bias_swiglu_impl(input_dtype, with_bias): + clamp_value = 10.0 + + if input_dtype == torch.float32: + tols = dict(rtol=1.0e-6, atol=1.0e-6) + elif input_dtype == torch.bfloat16: + tols = dict(rtol=2.0e-2, atol=1.0e-3) + else: + raise ValueError(f"Invalid input dtype: {input_dtype}") + + x = (torch.randn(16, 64, dtype=input_dtype, device="cuda") * 5.0).requires_grad_(True) + bias = ( + torch.randn(64, dtype=input_dtype, device="cuda").requires_grad_(True) + if with_bias + else None + ) + bwd_input = torch.randn(16, 32, dtype=input_dtype, device="cuda") + + x_fp32 = x.to(torch.float32) + x_effective = x_fp32 + bias.to(torch.float32) if with_bias else x_fp32 + y_1, y_2 = torch.chunk(x_effective, 2, -1) + y = ( + F.silu(y_1.clamp(min=None, max=clamp_value)) * y_2.clamp(min=-clamp_value, max=clamp_value) + ).to(input_dtype) + y.backward(bwd_input) + + x_fused = x.detach().clone().requires_grad_(True) + bias_fused = bias.detach().clone().requires_grad_(True) if with_bias else None + y_fused = bias_swiglu_impl(x_fused, bias_fused, clamp_value=clamp_value) + y_fused.backward(bwd_input.detach().clone()) + + assert y_fused.dtype == y.dtype + assert torch.allclose(y, y_fused, **tols) + assert x_fused.grad.dtype == x.grad.dtype + assert torch.allclose(x.grad, x_fused.grad, **tols) + if with_bias: + assert bias_fused.grad.dtype == bias.grad.dtype + bias_grad_cos = F.cosine_similarity( + bias.grad.flatten().float().unsqueeze(0), bias_fused.grad.flatten().float().unsqueeze(0) + ).item() + assert bias_grad_cos > 0.999, f"bias.grad cosine similarity = {bias_grad_cos:.6f}" + + +@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_bias_swiglu_impl_clamp_none_matches_unclamped(input_dtype, with_bias): + if input_dtype == torch.float32: + tols = dict(rtol=1.0e-6, atol=1.0e-6) + else: + tols = dict(rtol=2.0e-2, atol=1.0e-3) + + x = torch.randn(16, 64, dtype=input_dtype, device="cuda").requires_grad_(True) + bias = ( + torch.randn(64, dtype=input_dtype, device="cuda").requires_grad_(True) + if with_bias + else None + ) + bwd_input = torch.randn(16, 32, dtype=input_dtype, device="cuda") + + y = bias_swiglu_impl(x, bias) + y.backward(bwd_input) + + x_explicit = x.detach().clone().requires_grad_(True) + bias_explicit = bias.detach().clone().requires_grad_(True) if with_bias else None + y_explicit = bias_swiglu_impl(x_explicit, bias_explicit, clamp_value=None) + y_explicit.backward(bwd_input.detach().clone()) + + assert torch.allclose(y, y_explicit, **tols) + assert torch.allclose(x.grad, x_explicit.grad, **tols) + if with_bias: + bias_grad_cos = F.cosine_similarity( + bias.grad.flatten().float().unsqueeze(0), + bias_explicit.grad.flatten().float().unsqueeze(0), + ).item() + assert bias_grad_cos > 0.999, f"bias.grad cosine similarity = {bias_grad_cos:.6f}" diff --git a/tests/unit_tests/models/test_hybrid_hash_routing.py b/tests/unit_tests/models/test_hybrid_hash_routing.py new file mode 100644 index 00000000000..afc489bf1af --- /dev/null +++ b/tests/unit_tests/models/test_hybrid_hash_routing.py @@ -0,0 +1,387 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch + +from megatron.core import recompute as recompute_module +from megatron.core.models.hybrid.hybrid_block import HybridStack, HybridStackSubmodules +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols as LayerSymbols +from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.transformer.moe import router as router_module +from megatron.core.transformer.moe.router import TopKRouter +from megatron.core.transformer.multi_token_prediction import ( + MultiTokenPredictionLayer, + MultiTokenPredictionLayerSubmodules, +) +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.transformer.transformer_layer import TransformerLayer + + +class RecordingTransformerLayer(TransformerLayer): + """Minimal TransformerLayer stand-in that records its call signature.""" + + def __init__(self, layer_number=1, cuda_graph_impl="none"): + torch.nn.Module.__init__(self) + self.layer_number = layer_number + self.config = SimpleNamespace(cuda_graph_impl=cuda_graph_impl) + self.calls = [] + self.used_cuda_graph = False + if cuda_graph_impl == "local": + self.cudagraph_manager = self._run_local_cuda_graph + if cuda_graph_impl == "transformer_engine": + self.cuda_graphs = [object()] + + def forward(self, hidden_states, **kwargs): + self.calls.append(kwargs) + return hidden_states, None + + def _te_cuda_graph_replay(self, *args, **kwargs): + self.used_cuda_graph = True + return self.forward(*args, **kwargs) + + def _run_local_cuda_graph(self, module, args, kwargs): + assert module is self + self.used_cuda_graph = True + return self.forward(*args, **kwargs) + + +class RecordingNonTransformerLayer(torch.nn.Module): + """Mamba-like layer whose signature deliberately excludes input_ids.""" + + def __init__(self, layer_number=2): + super().__init__() + self.layer_number = layer_number + self.num_calls = 0 + + def forward( + self, + hidden_states, + attention_mask=None, + inference_context=None, + rotary_pos_emb=None, + *, + packed_seq_params=None, + ): + self.num_calls += 1 + return hidden_states + + +class RecordingDecoder: + def __init__(self): + self.kwargs = None + + def __call__(self, **kwargs): + self.kwargs = kwargs + return kwargs['hidden_states'] + + +def make_stack(layers, **config_overrides): + config = { + 'cuda_graph_impl': "none", + 'flash_decode': False, + 'fp8': False, + 'fp8_recipe': None, + 'fp4': None, + 'enable_hyper_connections': False, + 'recompute_granularity': None, + 'recompute_method': None, + 'recompute_num_layers': None, + 'distribute_saved_activations': False, + } + config.update(config_overrides) + stack = SimpleNamespace( + config=SimpleNamespace(**config), + pre_process=True, + post_process=False, + post_layer_norm=False, + input_tensor=None, + layers=layers, + num_layers_per_pipeline_rank=len(layers), + training=True, + _mhc_block_end_plan=None, + ) + stack._build_mhc_recompute_layer_plan = lambda _enabled: ( + [None] * len(layers), + [False] * len(layers), + ) + stack._finalize_mhc_recompute_layer = lambda **_kwargs: None + return stack + + +def run_stack(stack, input_ids): + hidden_states = torch.randn(4, 2, 8, requires_grad=True) + output = HybridStack.forward( + stack, hidden_states=hidden_states, attention_mask=None, input_ids=input_ids + ) + assert output.shape == hidden_states.shape + return output + + +@pytest.mark.parametrize("recompute_granularity", [None, "selective"]) +def test_hybrid_stack_forwards_input_ids_only_to_transformer_layers(recompute_granularity): + transformer_layer = RecordingTransformerLayer() + non_transformer_layer = RecordingNonTransformerLayer() + stack = make_stack( + [transformer_layer, non_transformer_layer], recompute_granularity=recompute_granularity + ) + input_ids = torch.arange(8).reshape(2, 4) + + run_stack(stack, input_ids) + + assert transformer_layer.calls[0]['input_ids'] is input_ids + assert non_transformer_layer.num_calls == 1 + + +def test_hybrid_stack_omits_input_ids_keyword_when_not_provided(): + transformer_layer = RecordingTransformerLayer() + stack = make_stack([transformer_layer]) + + run_stack(stack, input_ids=None) + + assert 'input_ids' not in transformer_layer.calls[0] + + +def test_hybrid_stack_full_recompute_preserves_ids_and_non_transformer_signature(monkeypatch): + monkeypatch.setattr( + recompute_module.tensor_parallel, + "checkpoint", + lambda function, _distribute_saved_activations, *args: function(*args), + ) + transformer_layer = RecordingTransformerLayer() + non_transformer_layer = RecordingNonTransformerLayer() + stack = make_stack( + [transformer_layer, non_transformer_layer], + recompute_granularity="full", + recompute_method="uniform", + recompute_num_layers=2, + ) + input_ids = torch.arange(8).reshape(2, 4) + + run_stack(stack, input_ids) + + assert transformer_layer.calls[0]['input_ids'] is input_ids + assert non_transformer_layer.num_calls == 1 + + +@pytest.mark.parametrize("cuda_graph_impl", ["local", "transformer_engine"]) +def test_hybrid_stack_preserves_hash_ids_in_cuda_graph_signature(cuda_graph_impl): + transformer_layer = RecordingTransformerLayer(cuda_graph_impl=cuda_graph_impl) + stack = make_stack([transformer_layer]) + input_ids = torch.arange(8).reshape(2, 4) + + run_stack(stack, input_ids) + + assert transformer_layer.used_cuda_graph + assert transformer_layer.calls[0]['input_ids'] is input_ids + + +@pytest.mark.parametrize("moe_n_hash_layers,expects_input_ids", [(0, False), (1, True)]) +def test_hybrid_model_passes_ids_to_decoder_only_for_hash_routing( + moe_n_hash_layers, expects_input_ids +): + decoder = RecordingDecoder() + config = SimpleNamespace( + fine_grained_activation_offloading=False, + moe_paged_stash=False, + moe_n_hash_layers=moe_n_hash_layers, + actual_vocab_size=128, + ) + model = SimpleNamespace( + config=config, + decoder=decoder, + position_embedding_type='none', + pre_process=True, + post_process=False, + share_embeddings_and_output_weights=False, + mtp_process=False, + vocab_size=128, + ) + input_ids = torch.arange(8).reshape(2, 4) + hidden_states = torch.randn(4, 2, 8) + + output = HybridModel.forward( + model, + input_ids=input_ids, + position_ids=torch.arange(4).repeat(2, 1), + attention_mask=None, + decoder_input=hidden_states, + ) + + assert output is hidden_states + assert config.actual_vocab_size == 128 + assert ('input_ids' in decoder.kwargs) is expects_input_ids + if expects_input_ids: + assert decoder.kwargs['input_ids'] is input_ids + + +def test_hybrid_hash_moe_pp_does_not_require_explicit_pipeline_layout(): + config = TransformerConfig( + num_layers=4, + hidden_size=64, + num_attention_heads=4, + use_cpu_initialization=True, + pipeline_model_parallel_size=2, + num_moe_experts=4, + moe_n_hash_layers=3, + actual_vocab_size=128, + is_hybrid_model=True, + ) + + assert config.pipeline_model_parallel_layout is None + + with pytest.raises( + AssertionError, + match="pipeline_model_parallel_layout must be set", + ): + TransformerConfig( + num_layers=4, + hidden_size=64, + num_attention_heads=4, + use_cpu_initialization=True, + pipeline_model_parallel_size=2, + num_moe_experts=4, + moe_n_hash_layers=3, + actual_vocab_size=128, + is_hybrid_model=False, + ) + + +def test_hybrid_stack_marks_mtp_moe_and_propagates_mtp_depth(monkeypatch): + import megatron.core.models.hybrid.hybrid_block as hybrid_block_module + + captured_build_kwargs = {} + + class _MtpMoEStub(torch.nn.Module): + def __init__(self, layer_number, is_mtp_layer): + super().__init__() + self.layer_number = layer_number + self.router = SimpleNamespace(is_mtp_layer=is_mtp_layer) + + def fake_build_module(_spec, **kwargs): + captured_build_kwargs.update(kwargs) + return _MtpMoEStub( + layer_number=kwargs["layer_number"], + is_mtp_layer=kwargs["is_mtp_layer"], + ) + + monkeypatch.setattr(hybrid_block_module, "build_module", fake_build_module) + config = SimpleNamespace( + fp8=False, + fp4=None, + enable_hyper_connections=False, + cuda_graph_impl="none", + ) + + stack = HybridStack( + config=config, + submodules=HybridStackSubmodules(moe_layer=object()), + layer_type_list=[LayerSymbols.MOE], + post_process=False, + pg_collection=SimpleNamespace(pp=object(), tp=object()), + is_mtp_layer=True, + mtp_layer_number=2, + ) + + assert captured_build_kwargs["is_mtp_layer"] is True + assert stack.is_mtp_layer is True + assert stack.mtp_layer_number == 2 + assert stack.layers[0].router.is_mtp_layer is True + assert stack.layers[0].router.mtp_layer_number == 2 + + +def test_mtp_layer_passes_its_depth_to_nested_hybrid_stack(monkeypatch): + import megatron.core.models.hybrid.hybrid_block as hybrid_block_module + import megatron.core.models.hybrid.hybrid_layer_allocation as allocation_module + import megatron.core.transformer.multi_token_prediction as mtp_module + + captured_stack_kwargs = {} + + class _IdentityNorm(torch.nn.Module): + def __init__(self, **_kwargs): + super().__init__() + + def forward(self, hidden_states): + return hidden_states + + class _RecordingHybridStack(torch.nn.Module): + def __init__(self, **kwargs): + super().__init__() + captured_stack_kwargs.update(kwargs) + + monkeypatch.setattr( + hybrid_block_module, "HybridStack", _RecordingHybridStack + ) + monkeypatch.setattr( + allocation_module, + "validate_segment_layers", + lambda _pattern: [LayerSymbols.MOE], + ) + monkeypatch.setattr( + mtp_module, + "build_module", + lambda *_args, **_kwargs: torch.nn.Identity(), + ) + + config = SimpleNamespace( + enable_hyper_connections=False, + sequence_parallel=False, + pipeline_model_parallel_size=1, + pipeline_model_parallel_layout=None, + hidden_size=8, + layernorm_epsilon=1e-5, + init_method=lambda tensor: tensor, + mtp_num_layers=2, + ) + submodules = MultiTokenPredictionLayerSubmodules( + enorm=_IdentityNorm, + hnorm=_IdentityNorm, + layer_norm=_IdentityNorm, + eh_proj=object(), + mtp_model_layer=None, + ) + + layer = MultiTokenPredictionLayer( + config=config, + submodules=submodules, + layer_number=2, + pg_collection=SimpleNamespace(cp=None, tp=None), + mtp_layer_pattern="E", + hybrid_submodules=HybridStackSubmodules(), + ) + + assert layer.layer_number == 2 + assert captured_stack_kwargs["is_mtp_layer"] is True + assert captured_stack_kwargs["mtp_layer_number"] == 2 + + +def test_hybrid_mtp_aux_metric_uses_enclosing_depth_slot(): + """An internal `/WE` MoE logs to its MTP depth, not its Hybrid sublayer number.""" + router = TopKRouter.__new__(TopKRouter) + torch.nn.Module.__init__(router) + router.config = SimpleNamespace( + mtp_num_layers=1, + mtp_use_repeated_layer=False, + num_layers=86, + ) + router.is_mtp_layer = True + router.layer_number = 2 + router.mtp_layer_number = 1 + router.calculate_per_token_loss = False + + activation = torch.ones(2) + tracker = mock.MagicMock() + with mock.patch.object(router_module, "get_moe_metrics_tracker", return_value=tracker): + router.attach_and_log_load_balancing_loss( + activation, + aux_loss_coeff=0.1, + aux_loss=torch.tensor(0.5), + aux_loss_name="seq_load_balancing_loss", + reduce_group=mock.sentinel.reduce_group, + ) + + record_args = tracker.record.call_args.args + assert record_args[2] == 87 + assert record_args[3] == 87 diff --git a/tests/unit_tests/models/test_hybrid_mhc.py b/tests/unit_tests/models/test_hybrid_mhc.py new file mode 100644 index 00000000000..ce87d8ed72b --- /dev/null +++ b/tests/unit_tests/models/test_hybrid_mhc.py @@ -0,0 +1,449 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import pytest +import torch + +from megatron.core.models.hybrid.hybrid_block import ( + HybridStack, + HybridStackSubmodules, + HyperConnectionHybridLayer, +) +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer import TransformerConfig +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.moe.moe_layer import MoELayer +from megatron.core.transformer.multi_token_prediction import MultiTokenPredictionLayer +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_layer import TransformerLayer +from tests.unit_tests.test_utilities import Utils + + +class _DummyHybridLayer(MegatronModule): + """Same-shape residual layer used to isolate HybridStack mHC plumbing.""" + + def __init__(self, config: TransformerConfig, layer_number: int, **_kwargs): + super().__init__(config=config) + self.layer_number = layer_number + self.proj = torch.nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self.seen_hidden_shapes = [] + + def forward( + self, + hidden_states, + attention_mask=None, + inference_context=None, + packed_seq_params=None, + **_kwargs, + ): + self.seen_hidden_shapes.append(tuple(hidden_states.shape)) + return hidden_states + 0.125 * self.proj(hidden_states) + + +class _StubTransformerLayer(TransformerLayer): + """Minimal TransformerLayer that exercises only the mHC wrapper guard.""" + + def __init__(self, config: TransformerConfig): + torch.nn.Module.__init__(self) + self.config = config + self.layer_number = 1 + + def _forward_attention(self, *args, **kwargs): + hidden_states = kwargs.get("hidden_states", args[0] if args else None) + return hidden_states, None + + def _forward_mlp(self, hidden_states, *_args, **_kwargs): + return hidden_states + + +def _get_pg_collection() -> ProcessGroupCollection: + return ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'pp', 'cp']) + + +def _get_dummy_submodules() -> HybridStackSubmodules: + layer_spec = ModuleSpec(module=_DummyHybridLayer) + return HybridStackSubmodules( + mamba_layer=layer_spec, + gdn_layer=layer_spec, + attention_layer=layer_spec, + dsa_layer=layer_spec, + mlp_layer=layer_spec, + moe_layer=layer_spec, + ) + + +def _get_dummy_stack_spec() -> ModuleSpec: + return ModuleSpec( + module=HybridStack, params={"post_layer_norm": False}, submodules=_get_dummy_submodules() + ) + + +def _get_config(num_layers: int, **kwargs) -> TransformerConfig: + return TransformerConfig( + num_layers=num_layers, + hidden_size=32, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=2, + hidden_dropout=0.0, + mhc_sinkhorn_iterations=3, + **kwargs, + ) + + +def _get_stack( + config: TransformerConfig, + num_local_layers: int, + *, + pre_process: bool = True, + post_process: bool = True, + pp_layer_offset: int = 0, +) -> HybridStack: + return HybridStack( + config=config, + submodules=_get_dummy_submodules(), + pre_process=pre_process, + post_process=post_process, + post_layer_norm=False, + layer_type_list=[Symbols.MAMBA] * num_local_layers, + pp_layer_offset=pp_layer_offset, + pg_collection=_get_pg_collection(), + ) + + +def test_mhc_mtp_requires_hybrid_contract(): + config = _get_config(num_layers=1, mtp_num_layers=1) + + with pytest.raises(ValueError, match="requires the HybridModel MTP contract"): + MultiTokenPredictionLayer( + config=config, + submodules=object(), + layer_number=1, + pg_collection=None, + ) + + +@pytest.mark.internal +class TestHybridStackMHC: + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def test_constructor_and_sharded_state(self): + config = _get_config(num_layers=3) + stack = _get_stack(config, num_local_layers=3) + + assert all(isinstance(layer, HyperConnectionHybridLayer) for layer in stack.layers) + assert stack.hc_head_fn.shape == ( + config.num_residual_streams, + config.hidden_size * config.num_residual_streams, + ) + state = stack.sharded_state_dict(prefix="decoder.", metadata={}) + for name in ("hc_head_fn", "hc_head_base", "hc_head_scale"): + assert f"decoder.{name}" in state + + @pytest.mark.parametrize( + "recompute_kwargs", + [ + {}, + { + "recompute_granularity": "selective", + "recompute_modules": ["core_attn", "mhc"], + "mhc_recompute_layer_num": 2, + }, + { + "recompute_granularity": "full", + "recompute_method": "uniform", + "recompute_num_layers": 1, + }, + ], + ids=["none", "selective_mhc", "full_uniform"], + ) + def test_forward_backward(self, recompute_kwargs): + config = _get_config(num_layers=3, **recompute_kwargs) + stack = _get_stack(config, num_local_layers=3).cuda() + hidden_states = torch.randn(8, 2, config.hidden_size, device="cuda", requires_grad=True) + + output = stack(hidden_states, attention_mask=None) + + assert output.shape == hidden_states.shape + assert torch.isfinite(output).all() + output.float().sum().backward() + assert hidden_states.grad is not None + for layer in stack.layers: + assert layer.inner_layer.proj.weight.grad is not None + assert layer.hyper_connection.mapping_proj.weight.grad is not None + assert all( + shape == (8, 2, config.hidden_size) + for shape in layer.inner_layer.seen_hidden_shapes + ) + for name in ("hc_head_fn", "hc_head_base", "hc_head_scale"): + assert getattr(stack, name).grad is not None + + def test_full_recompute_forwards_input_ids(self, monkeypatch): + config = _get_config( + num_layers=2, + recompute_granularity="full", + recompute_method="uniform", + recompute_num_layers=1, + ) + stack = _get_stack(config, num_local_layers=2).cuda() + hidden_states = torch.randn(8, 2, config.hidden_size, device="cuda", requires_grad=True) + input_ids = torch.arange(8, device="cuda").repeat(2, 1) + seen_input_ids = [] + original = HyperConnectionHybridLayer._call_inner_transformer_layer_without_local_bda + + def record_input_ids( + layer, + hidden_states, + attention_mask, + inference_context, + rotary_pos_emb, + sequence_len_offset, + packed_seq_params, + padding_mask, + input_ids=None, + ): + seen_input_ids.append(input_ids) + return original( + layer, + hidden_states, + attention_mask, + inference_context, + rotary_pos_emb, + sequence_len_offset, + packed_seq_params, + padding_mask, + input_ids, + ) + + monkeypatch.setattr( + HyperConnectionHybridLayer, + "_call_inner_transformer_layer_without_local_bda", + record_input_ids, + ) + + output = stack(hidden_states, attention_mask=None, input_ids=input_ids) + output.float().sum().backward() + + assert len(seen_input_ids) == 4 + assert all( + captured is not None and torch.equal(captured, input_ids) for captured in seen_input_ids + ) + + def test_fused_bf16_forward_backward(self): + config = _get_config( + num_layers=2, bf16=True, params_dtype=torch.bfloat16, use_fused_mhc=True + ) + stack = _get_stack(config, num_local_layers=2).cuda().bfloat16() + hidden_states = torch.randn( + 8, 2, config.hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + + output = stack(hidden_states, attention_mask=None) + + assert output.shape == hidden_states.shape + assert output.dtype == torch.bfloat16 + assert torch.isfinite(output).all() + output.float().sum().backward() + assert hidden_states.grad is not None + assert all( + layer.hyper_connection.mapping_proj.weight.grad is not None for layer in stack.layers + ) + + def test_pipeline_boundary_shapes(self): + config = _get_config(num_layers=2) + first_stage = _get_stack( + config, num_local_layers=1, pre_process=True, post_process=False + ).cuda() + last_stage = _get_stack( + config, num_local_layers=1, pre_process=False, post_process=True, pp_layer_offset=1 + ).cuda() + hidden_states = torch.randn(8, 2, config.hidden_size, device="cuda") + + pipeline_hidden = first_stage(hidden_states, attention_mask=None) + assert pipeline_hidden.shape == (8, 2, config.hidden_size * config.num_residual_streams) + + last_stage.set_input_tensor(pipeline_hidden.detach()) + output = last_stage(hidden_states, attention_mask=None) + assert output.shape == hidden_states.shape + + def test_real_attention_mlp_forward_backward(self): + config = _get_config(num_layers=2) + stack = HybridStack( + config=config, + submodules=hybrid_stack_spec.submodules, + post_layer_norm=False, + layer_type_list=[Symbols.ATTENTION, Symbols.MLP], + pp_layer_offset=0, + pg_collection=_get_pg_collection(), + ).cuda() + hidden_states = torch.randn(8, 2, config.hidden_size, device="cuda", requires_grad=True) + + output = stack(hidden_states, attention_mask=None) + + assert output.shape == hidden_states.shape + assert all(isinstance(layer, HyperConnectionHybridLayer) for layer in stack.layers) + assert all(isinstance(layer.inner_layer, TransformerLayer) for layer in stack.layers) + output.float().sum().backward() + assert hidden_states.grad is not None + assert all( + layer.hyper_connection.mapping_proj.weight.grad is not None for layer in stack.layers + ) + + def test_real_moe_raw_branch_forward_backward(self): + config = _get_config( + num_layers=1, + num_moe_experts=2, + moe_ffn_hidden_size=64, + moe_grouped_gemm=True, + add_bias_linear=False, + ) + pg_collection = ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['tp', 'pp', 'cp', 'ep', 'expt_tp', 'tp_ep', 'expt_dp'] + ) + stack = HybridStack( + config=config, + submodules=hybrid_stack_spec.submodules, + post_layer_norm=False, + layer_type_list=[Symbols.MOE], + pp_layer_offset=0, + pg_collection=pg_collection, + ).cuda() + hidden_states = torch.randn(8, 2, config.hidden_size, device="cuda", requires_grad=True) + + output = stack(hidden_states, attention_mask=None) + + wrapped_layer = stack.layers[0] + assert isinstance(wrapped_layer, HyperConnectionHybridLayer) + assert isinstance(wrapped_layer.inner_layer.mlp, MoELayer) + assert output.shape == hidden_states.shape + output.float().sum().backward() + assert hidden_states.grad is not None + assert wrapped_layer.hyper_connection.mapping_proj.weight.grad is not None + assert any(param.grad is not None for param in wrapped_layer.inner_layer.mlp.parameters()) + + def test_hybrid_model_forward_backward(self): + config = _get_config(num_layers=3) + model = HybridModel( + config=config, + hybrid_stack_spec=_get_dummy_stack_spec(), + vocab_size=64, + max_sequence_length=8, + hybrid_layer_pattern="M*-", + parallel_output=False, + ).cuda() + input_ids = torch.arange(8, dtype=torch.int64, device="cuda").repeat((2, 1)) + position_ids = torch.arange(8, dtype=torch.int64, device="cuda").repeat((2, 1)) + + logits = model(input_ids=input_ids, position_ids=position_ids, attention_mask=None) + + assert logits.shape == (2, 8, model.vocab_size) + assert torch.isfinite(logits).all() + logits.float().mean().backward() + assert all(layer.inner_layer.proj.weight.grad is not None for layer in model.decoder.layers) + assert all( + layer.hyper_connection.mapping_proj.weight.grad is not None + for layer in model.decoder.layers + ) + + def test_hybrid_model_mtp_forward_backward(self): + config = _get_config(num_layers=1, mtp_num_layers=1, mtp_loss_scaling_factor=0.1) + model = HybridModel( + config=config, + hybrid_stack_spec=hybrid_stack_spec, + vocab_size=64, + max_sequence_length=8, + hybrid_layer_pattern="-/-", + parallel_output=True, + ).cuda() + input_ids = torch.arange(8, dtype=torch.int64, device="cuda").repeat((2, 1)) + position_ids = torch.arange(8, dtype=torch.int64, device="cuda").repeat((2, 1)) + + logits = model(input_ids=input_ids, position_ids=position_ids, attention_mask=None) + + assert logits.shape == (2, 8, model.vocab_size) + assert torch.isfinite(logits).all() + assert not any("mtp_model_layer.hc_head_" in name for name, _ in model.named_parameters()) + logits.float().mean().backward() + mtp_params = [param for name, param in model.named_parameters() if name.startswith("mtp.")] + assert mtp_params + assert all(param.grad is not None for param in mtp_params) + + def test_recompute_plan(self): + config = _get_config( + num_layers=3, + recompute_granularity="selective", + recompute_modules=["core_attn", "mhc"], + mhc_recompute_layer_num=2, + ) + stack = _get_stack(config, num_local_layers=3) + + managers, block_ends = stack._build_mhc_recompute_layer_plan(True) + + assert block_ends == [False, True, True] + assert managers[0] is managers[1] + assert managers[1] is not managers[2] + + def test_boundary_bda_skips_recompute_manager(self, monkeypatch): + config = _get_config(num_layers=1) + layer = HyperConnectionHybridLayer( + config=config, layer=_DummyHybridLayer(config, layer_number=1) + ) + hidden_states = torch.randn( + 4, 2, config.hidden_size * config.num_residual_streams, requires_grad=True + ) + manager = type("_FakeManager", (), {})() + manager.is_last_layer_in_recompute_block = True + seen_managers = [] + + def fake_hyper_connection_forward( + hidden_states, mhc_recompute_manager=None, return_residual=False + ): + assert mhc_recompute_manager is manager + assert return_residual + sequence_length, batch_size, _ = hidden_states.shape + n = config.num_residual_streams + hidden_size = config.hidden_size + aggregated = hidden_states.view(sequence_length, batch_size, n, hidden_size).mean(dim=2) + h_res = torch.empty(sequence_length, batch_size, n, n) + h_post = torch.empty(sequence_length, batch_size, n) + return aggregated, h_res, h_post, hidden_states + + def fake_bda( + h_res, residual, h_post, output_with_bias, dropout_prob, training, fused, manager=None + ): + seen_managers.append(manager) + return residual + + monkeypatch.setattr(layer.hyper_connection, "forward", fake_hyper_connection_forward) + monkeypatch.setattr(layer.hyper_connection, "fused_h_res_h_post_bda", fake_bda) + + output, _ = layer(hidden_states, attention_mask=None, mhc_recompute_manager=manager) + assert output is hidden_states + assert seen_managers == [None] + + manager.is_last_layer_in_recompute_block = False + layer(hidden_states, attention_mask=None, mhc_recompute_manager=manager) + assert seen_managers[-1] is manager + + def test_transformer_layer_wrapper_escape_hatch(self): + config = _get_config(num_layers=1) + layer = _StubTransformerLayer(config) + hidden_states = torch.randn(4, 2, config.hidden_size) + + with pytest.raises(RuntimeError, match="must not be called directly"): + layer.forward(hidden_states=hidden_states, attention_mask=None) + + output, context = layer.forward( + hidden_states=hidden_states, attention_mask=None, _called_from_hybrid_mhc_wrapper=True + ) + assert output is hidden_states + assert context is None diff --git a/tests/unit_tests/models/test_hybrid_model.py b/tests/unit_tests/models/test_hybrid_model.py index ffc9fe41e99..10ec44c2ce5 100644 --- a/tests/unit_tests/models/test_hybrid_model.py +++ b/tests/unit_tests/models/test_hybrid_model.py @@ -49,6 +49,52 @@ def test_hybrid_logging_process_groups_are_paired(): _hybrid_logging_pg_kwargs(SimpleNamespace(tp=None, dp_cp=dp_cp_group)) +class _EchoHybridDecoder(torch.nn.Module): + def forward(self, hidden_states, **_kwargs): + return hidden_states + + +class _RecordingMTP(torch.nn.Module): + def __init__(self): + super().__init__() + self.padding_mask = None + + def forward(self, hidden_states, padding_mask=None, **_kwargs): + self.padding_mask = padding_mask + return hidden_states + + +def test_hybrid_model_forwards_padding_mask_to_mtp(): + model = HybridModel.__new__(HybridModel) + torch.nn.Module.__init__(model) + model.config = SimpleNamespace( + fine_grained_activation_offloading=False, + moe_paged_stash=False, + moe_n_hash_layers=0, + ) + model.position_embedding_type = "none" + model.decoder = _EchoHybridDecoder() + model.mtp = _RecordingMTP() + model.embedding = None + model.mtp_process = True + model.post_process = False + model.share_embeddings_and_output_weights = False + hidden_states = torch.randn(4, 2, 8) + padding_mask = torch.tensor( + [[False, False, True, True], [False, True, False, True]], dtype=torch.bool + ) + + model( + input_ids=torch.zeros((2, 4), dtype=torch.long), + position_ids=torch.arange(4).repeat(2, 1), + attention_mask=None, + decoder_input=hidden_states, + padding_mask=padding_mask, + ) + + assert model.mtp.padding_mask is padding_mask + + @pytest.mark.skipif( not is_torch_min_version("2.4.0"), reason="torch.distributed.init_device_mesh requires torch >= 2.4.0", diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index f7dc78ce9a2..d4a0993ba45 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -33,8 +33,10 @@ "activation_func": "megatron.core.activations.squared_relu", "activation_func_clamp_value": None, "activation_func_fp8_input_store": False, + "actual_vocab_size": 131072, "add_bias_linear": False, "add_qkv_bias": False, + "apply_dsa_kernel_fusion": True, "apply_query_key_layer_scaling": False, "apply_residual_connection_post_layernorm": False, "apply_rope_fusion": False, @@ -60,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, @@ -68,6 +71,10 @@ "cpu_offloading_weights": False, "cross_entropy_fusion_impl": "native", "cross_entropy_loss_fusion": True, + "csa_compress_ratios": None, + "csa_compress_rotary_base": 40000.0, + "csa_dense_mode": False, + "csa_window_size": 128, "cuda_graph_impl": "none", "cuda_graph_retain_backward_graph": False, "cuda_graph_modules": [], @@ -100,6 +107,7 @@ "embedding_init_method_std": 0.014, "enable_autocast": False, "enable_cuda_graph": False, + "enable_hyper_connections": False, "ep_overlap_early_attn_memory_release": False, "experimental_attention_variant": None, "experimental_attention_variant_loss_scale_func": None, @@ -167,6 +175,9 @@ "mamba_training_ssm_states_dtype": None, "masked_softmax_fusion": True, "memory_efficient_layer_norm": False, + "mhc_init_gating_factor": 0.01, + "mhc_recompute_layer_num": None, + "mhc_sinkhorn_iterations": 20, "microbatch_group_size_per_vp_stage": 1, "mlp_chunks_for_prefill": 1, "mlp_chunks_for_training": 1, @@ -189,6 +200,7 @@ "moe_latent_size": None, "moe_layer_freq": 1, "moe_layer_recompute": False, + "moe_n_hash_layers": 0, "moe_ncclep_static_shape": False, "moe_ncclep_use_symm_mem": False, "moe_pad_expert_input_to_capacity": False, @@ -253,6 +265,7 @@ "num_microbatches_with_partial_activation_checkpoints": None, "num_moe_experts": 128, "num_query_groups": 2, + "num_residual_streams": 4, "output_layer_init_method": {}, "overlap_moe_expert_parallel_comm": False, "overlap_p2p_comm": False, @@ -337,6 +350,8 @@ "use_transformer_engine_op_fuser": False, "moe_single_grouped_weight": False, "moe_single_grouped_bias": False, + "moe_hybridep_pad_uneven_dispatch_inputs": False, + "sequence_packing_scheduler": None, } # Fields to ignore entirely (ephemeral, environment-specific, very large). SKIP_FIELDS = set() diff --git a/tests/unit_tests/pipeline_parallel/test_mhc_tensor_shapes.py b/tests/unit_tests/pipeline_parallel/test_mhc_tensor_shapes.py new file mode 100644 index 00000000000..2f72710197a --- /dev/null +++ b/tests/unit_tests/pipeline_parallel/test_mhc_tensor_shapes.py @@ -0,0 +1,153 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from types import SimpleNamespace + +import pytest +import torch + +from megatron.core.pipeline_parallel.schedules import _get_pipeline_hidden_size, get_tensor_shapes +from megatron.core.transformer.transformer_config import TransformerConfig + + +class FakeProcessGroup: + """Small process-group stand-in for schedule shape tests.""" + + def __init__(self, rank: int, size: int): + self._rank = rank + self._size = size + + def rank(self) -> int: + return self._rank + + def size(self) -> int: + return self._size + + +def make_config(**overrides): + values = { + 'hidden_size': 64, + 'enable_hyper_connections': True, + 'num_residual_streams': 4, + 'sequence_parallel': False, + 'variable_seq_lengths': False, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def get_shapes(config, *, pp_rank=0, pp_size=2, is_recv=True, tp_size=1, cp_size=1): + return get_tensor_shapes( + seq_length=32, + micro_batch_size=2, + decoder_seq_length=None, + config=config, + tp_group=FakeProcessGroup(0, tp_size), + cp_group=FakeProcessGroup(0, cp_size), + pp_group=FakeProcessGroup(pp_rank, pp_size), + is_recv=is_recv, + ) + + +@pytest.mark.parametrize( + "pp_rank,is_recv,hidden_size", + [ + (0, True, 64), + (0, False, 256), + (1, True, 256), + (1, False, 256), + (2, True, 256), + (2, False, 256), + (3, True, 256), + (3, False, 64), + ], +) +def test_non_interleaved_mhc_uses_expanded_shape_between_stages(pp_rank, is_recv, hidden_size): + shapes = get_shapes(make_config(), pp_rank=pp_rank, pp_size=4, is_recv=is_recv) + + assert shapes == [(32, 2, hidden_size)] + + +@pytest.mark.parametrize("enable_hyper_connections", [False, True]) +@pytest.mark.parametrize("is_recv", [False, True]) +def test_single_pipeline_stage_keeps_model_hidden_size(enable_hyper_connections, is_recv): + config = make_config(enable_hyper_connections=enable_hyper_connections) + + assert get_shapes(config, pp_size=1, is_recv=is_recv) == [(32, 2, 64)] + + +def test_non_mhc_pipeline_keeps_model_hidden_size(): + config = make_config(enable_hyper_connections=False) + + for rank in range(4): + assert get_shapes(config, pp_rank=rank, pp_size=4, is_recv=True) == [(32, 2, 64)] + assert get_shapes(config, pp_rank=rank, pp_size=4, is_recv=False) == [(32, 2, 64)] + + +def test_legacy_get_tensor_shapes_call_without_pp_group_is_unchanged(): + config = make_config() + + shapes = get_tensor_shapes( + seq_length=32, + micro_batch_size=2, + decoder_seq_length=None, + config=config, + tp_group=FakeProcessGroup(0, 1), + cp_group=FakeProcessGroup(0, 1), + ) + + assert shapes == [(32, 2, 64)] + + +@pytest.mark.parametrize( + "enabled,pp_size,hidden_size", [(False, 1, 64), (False, 4, 64), (True, 1, 64), (True, 4, 256)] +) +def test_interleaved_pipeline_uses_one_shape_for_all_active_edges(enabled, pp_size, hidden_size): + config = make_config(enable_hyper_connections=enabled) + + assert _get_pipeline_hidden_size(config, pp_group=FakeProcessGroup(0, pp_size)) == hidden_size + + +def test_mhc_shape_preserves_context_and_sequence_parallel_scaling(): + config = make_config(sequence_parallel=True) + + shapes = get_shapes(config, pp_rank=1, pp_size=4, is_recv=True, tp_size=2, cp_size=4) + + assert shapes == [(4, 2, 256)] + + +def test_mhc_shape_uses_decoder_sequence_length(): + config = make_config() + + shapes = get_tensor_shapes( + seq_length=32, + micro_batch_size=2, + decoder_seq_length=48, + config=config, + tp_group=FakeProcessGroup(0, 1), + cp_group=FakeProcessGroup(0, 2), + pp_group=FakeProcessGroup(1, 4), + is_recv=True, + ) + + assert shapes == [(24, 2, 256)] + + +def test_variable_sequence_length_shape_is_unchanged(): + config = make_config(variable_seq_lengths=True) + + assert get_shapes(config, pp_rank=1, pp_size=4, is_recv=True) == [()] + + +def test_native_mhc_transformer_config_drives_pipeline_shape(): + config = TransformerConfig( + num_layers=8, + hidden_size=64, + num_attention_heads=4, + pipeline_model_parallel_size=2, + pipeline_dtype=torch.bfloat16, + enable_hyper_connections=True, + num_residual_streams=4, + use_cpu_initialization=True, + ) + + assert get_shapes(config, pp_rank=0, pp_size=2, is_recv=False) == [(32, 2, 256)] diff --git a/tests/unit_tests/pipeline_parallel/test_schedules.py b/tests/unit_tests/pipeline_parallel/test_schedules.py index 92db675d193..3d50763ea61 100644 --- a/tests/unit_tests/pipeline_parallel/test_schedules.py +++ b/tests/unit_tests/pipeline_parallel/test_schedules.py @@ -377,7 +377,8 @@ def test_dsa_indexer_loss_scale_accepts_dict_output_tensor(): ) -def test_dsa_indexer_loss_scale_defaults_from_variant_without_mutating_config(): +@pytest.mark.parametrize("variant", ["dsa", "dsv4_hybrid"]) +def test_indexer_loss_scale_defaults_from_variant_without_mutating_config(variant): from megatron.core.transformer.experimental_attention_variant.dsa import ( DSAIndexerLossAutoScaler, ) @@ -385,7 +386,7 @@ def test_dsa_indexer_loss_scale_defaults_from_variant_without_mutating_config(): config = SimpleNamespace( calculate_per_token_loss=True, experimental_attention_variant_loss_scale_func=None, - experimental_attention_variant='dsa', + experimental_attention_variant=variant, grad_scale_func=lambda tensor: tensor * 7.0, num_moe_experts=None, mtp_num_layers=None, diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py index dd6c85b2125..06f67bc7038 100644 --- a/tests/unit_tests/rl/test_rl_utils.py +++ b/tests/unit_tests/rl/test_rl_utils.py @@ -167,6 +167,64 @@ def create_test_args(self, **kwargs): set_global_variables(args, False) return args + def test_get_rl_packed_seq_params_for_cuda_graph_without_sequence_packing(self): + params = rl_utils.get_rl_packed_seq_params_for_cuda_graph( + seq_length=8, device=torch.device("cpu"), sequence_packing=False + ) + + assert params.qkv_format == 'thd' + assert params.max_seqlen_q == 8 + assert params.max_seqlen_kv == 8 + assert params.total_tokens == 8 + assert params.cu_seqlens_kv is params.cu_seqlens_q + assert params.cu_seqlens_q.dtype == torch.int32 + assert params.cu_seqlens_q.device.type == "cpu" + assert torch.equal(params.cu_seqlens_q, torch.tensor([0, 8], dtype=torch.int32)) + assert torch.equal(params.seq_idx, torch.zeros((1, 8), dtype=torch.int32)) + + def test_get_rl_packed_seq_params_for_cuda_graph_with_sequence_packing(self): + params = rl_utils.get_rl_packed_seq_params_for_cuda_graph( + seq_length=8, device=torch.device("cpu"), sequence_packing=True, max_sequences_per_bin=3 + ) + + expected_cu_seqlens = torch.tensor([0, 8, 8, 8, 8], dtype=torch.int32) + assert params.qkv_format == 'thd' + assert params.max_seqlen_q == 8 + assert params.max_seqlen_kv == 8 + assert params.total_tokens == 8 + assert params.cu_seqlens_kv is params.cu_seqlens_q + assert params.cu_seqlens_q.dtype == torch.int32 + assert params.cu_seqlens_q.device.type == "cpu" + assert params.cu_seqlens_q.shape == (5,) + assert torch.equal(params.cu_seqlens_q, expected_cu_seqlens) + assert torch.equal(params.seq_idx, torch.zeros((1, 8), dtype=torch.int32)) + + def test_get_rl_packed_seq_params_for_cuda_graph_requires_max_sequences_per_bin(self): + with pytest.raises(AssertionError, match="max_sequences_per_bin is required"): + rl_utils.get_rl_packed_seq_params_for_cuda_graph( + seq_length=8, device=torch.device("cpu"), sequence_packing=True + ) + + def test_get_rl_packed_seq_params_for_cuda_graph_edge_cases(self): + # Parametrize seq_length of 1 (single-token boundary condition) + params_single = rl_utils.get_rl_packed_seq_params_for_cuda_graph( + seq_length=1, device=torch.device("cpu"), sequence_packing=False + ) + assert params_single.max_seqlen_q == 1 + assert params_single.max_seqlen_kv == 1 + assert params_single.total_tokens == 1 + assert torch.equal(params_single.cu_seqlens_q, torch.tensor([0, 1], dtype=torch.int32)) + + # Parametrize sequence packing with max_sequences_per_bin > 1 (e.g. 4) + params_multi = rl_utils.get_rl_packed_seq_params_for_cuda_graph( + seq_length=8, device=torch.device("cpu"), sequence_packing=True, max_sequences_per_bin=4 + ) + assert params_multi.max_seqlen_q == 8 + assert params_multi.max_seqlen_kv == 8 + assert params_multi.total_tokens == 8 + assert params_multi.cu_seqlens_q.shape == (6,) + assert torch.equal(params_multi.cu_seqlens_q, torch.tensor([0, 8, 8, 8, 8, 8], dtype=torch.int32)) + def test_rl_granularity_defaults(self): args = self.create_test_args(perform_rl_step=True, grpo_prompts_per_step=8) diff --git a/tests/unit_tests/ssm/test_hybrid_layer_allocation.py b/tests/unit_tests/ssm/test_hybrid_layer_allocation.py index faa553216da..c0c1ac52552 100644 --- a/tests/unit_tests/ssm/test_hybrid_layer_allocation.py +++ b/tests/unit_tests/ssm/test_hybrid_layer_allocation.py @@ -19,6 +19,13 @@ ) +def expected_layer_counts(nonzero=None): + """Build an exact count map while keeping tests resilient to new zero-count symbols.""" + counts = {symbol: 0 for symbol in Symbols.VALID_LAYERS} + counts.update(nonzero or {}) + return counts + + @pytest.mark.internal class TestPatternFromRatios: @@ -78,6 +85,7 @@ def test_valid_patterns(self): ("GGG*GGG*", ['G', 'G', 'G', '*', 'G', 'G', 'G', '*']), ("GEGEGE*E", ['G', 'E', 'G', 'E', 'G', 'E', '*', 'E']), ("MDMD", ['M', 'D', 'M', 'D']), + ("WECEH", ['W', 'E', 'C', 'E', 'H']), ] for pattern, expected in test_cases: result = validate_segment_layers(pattern) @@ -102,6 +110,12 @@ def test_invalid_symbols_cause_failure(self): # Not allowed to have both standard Attention and MLA/DSA validate_segment_layers("MDM*-") + def test_dsv4_attention_symbols(self): + assert {Symbols.WINDOW, Symbols.CSA, Symbols.HCA} <= Symbols.MLA_ATTENTION + assert validate_segment_layers("WDCH") == ["W", "D", "C", "H"] + with pytest.raises(ValueError): + validate_segment_layers("W*C") + @pytest.mark.internal class TestGetHybridTotalLayerCount: @@ -305,74 +319,63 @@ def test_dataclass_equality(self): class TestGetHybridLayerCounts: def test_simple_pattern(self): - assert get_hybrid_layer_counts("M*M*") == {'*': 2, 'D': 0, 'G': 0, 'M': 2, '-': 0, 'E': 0} + assert get_hybrid_layer_counts("M*M*") == expected_layer_counts({'*': 2, 'M': 2}) def test_all_layer_types(self): # Not allowed to have both standard Attention and MLA/DSA, so we do separate asserts. - assert get_hybrid_layer_counts("MG*-E") == {'*': 1, 'D': 0, 'G': 1, 'M': 1, '-': 1, 'E': 1} - assert get_hybrid_layer_counts("MGD-E") == {'*': 0, 'D': 1, 'G': 1, 'M': 1, '-': 1, 'E': 1} + assert get_hybrid_layer_counts("MG*-E") == expected_layer_counts( + {'*': 1, 'G': 1, 'M': 1, '-': 1, 'E': 1} + ) + assert get_hybrid_layer_counts("MGD-E") == expected_layer_counts( + {'D': 1, 'G': 1, 'M': 1, '-': 1, 'E': 1} + ) + assert get_hybrid_layer_counts("WCHD") == expected_layer_counts( + {'W': 1, 'C': 1, 'H': 1, 'D': 1} + ) def test_with_pipes(self): # Pipes should be skipped in counting - assert get_hybrid_layer_counts("M*|M*") == {'*': 2, 'D': 0, 'G': 0, 'M': 2, '-': 0, 'E': 0} - assert get_hybrid_layer_counts("M-M-|M-M*-") == { - '*': 1, - 'D': 0, - 'G': 0, - 'M': 4, - '-': 4, - 'E': 0, - } + assert get_hybrid_layer_counts("M*|M*") == expected_layer_counts({'*': 2, 'M': 2}) + assert get_hybrid_layer_counts("M-M-|M-M*-") == expected_layer_counts( + {'*': 1, 'M': 4, '-': 4} + ) def test_with_mtp(self): # MTP pattern "MM" repeated 2 depths -> 4 extra mamba layers - assert get_hybrid_layer_counts("M*M*/MM/MM") == { - '*': 2, - 'D': 0, - 'G': 0, - 'M': 6, - '-': 0, - 'E': 0, - } + assert get_hybrid_layer_counts("M*M*/MM/MM") == expected_layer_counts( + {'*': 2, 'M': 6} + ) def test_with_pipes_and_mtp(self): # Main: M-M-|M-M*- -> 1 attn, 4 mamba, 4 mlp # MTP: MM x 2 depths -> +4 mamba - assert get_hybrid_layer_counts("M-M-|M-M*-/MM/MM") == { - '*': 1, - 'D': 0, - 'G': 0, - 'M': 8, - '-': 4, - 'E': 0, - } + assert get_hybrid_layer_counts("M-M-|M-M*-/MM/MM") == expected_layer_counts( + {'*': 1, 'M': 8, '-': 4} + ) def test_moe_pattern(self): - assert get_hybrid_layer_counts("MEME") == {'*': 0, 'D': 0, 'G': 0, 'M': 2, '-': 0, 'E': 2} + assert get_hybrid_layer_counts("MEME") == expected_layer_counts({'M': 2, 'E': 2}) def test_mtp_with_attention(self): # MTP pattern "*M" repeated 3 depths -> 3 attn + 3 mamba from MTP - assert get_hybrid_layer_counts("MMMM/*M/*M/*M") == { - '*': 3, - 'D': 0, - 'G': 0, - 'M': 7, - '-': 0, - 'E': 0, - } + assert get_hybrid_layer_counts("MMMM/*M/*M/*M") == expected_layer_counts( + {'*': 3, 'M': 7} + ) def test_gdn_pattern(self): - assert get_hybrid_layer_counts("GMGM") == {'*': 0, 'D': 0, 'G': 2, 'M': 2, '-': 0, 'E': 0} + assert get_hybrid_layer_counts("GMGM") == expected_layer_counts({'G': 2, 'M': 2}) def test_gdn_hybrid_pattern(self): # GDN + Mamba + Attention - assert get_hybrid_layer_counts("G*GM*") == {'*': 2, 'D': 0, 'G': 2, 'M': 1, '-': 0, 'E': 0} + assert get_hybrid_layer_counts("G*GM*") == expected_layer_counts( + {'*': 2, 'G': 2, 'M': 1} + ) def test_dsa_pattern(self): - assert get_hybrid_layer_counts("DMDM") == {'*': 0, 'D': 2, 'G': 0, 'M': 2, '-': 0, 'E': 0} + assert get_hybrid_layer_counts("DMDM") == expected_layer_counts({'D': 2, 'M': 2}) def test_empty_pattern(self): - assert get_hybrid_layer_counts("") == {'*': 0, 'D': 0, 'G': 0, 'M': 0, '-': 0, 'E': 0} + assert get_hybrid_layer_counts("") == expected_layer_counts() @pytest.mark.internal @@ -655,7 +658,7 @@ def test_standard_layer_types(self): """Standard symbols each produce a single-entry map at local index 0.""" maps = get_layer_maps_from_layer_type_list(["*", "M", "-", "E"]) # We always get all symbols returned, not only those contained in the pattern. - assert len(maps) == 6 + assert len(maps) == len(Symbols.VALID_LAYERS) attention_map, mamba_map, mlp_map, moe_map = operator.itemgetter( Symbols.ATTENTION, Symbols.MAMBA, Symbols.MLP, Symbols.MOE )(maps) diff --git a/tests/unit_tests/test_argument_utils.py b/tests/unit_tests/test_argument_utils.py index 7c0b30d3d56..e50619455a2 100644 --- a/tests/unit_tests/test_argument_utils.py +++ b/tests/unit_tests/test_argument_utils.py @@ -13,6 +13,7 @@ from megatron.training.argument_utils import ( ArgumentGroupFactory, TypeInferenceError, + _normalize_dsv4_hybrid_csa_compress_ratios, pretrain_cfg_container_from_args, ) from megatron.training.config import PretrainConfigContainer @@ -83,6 +84,56 @@ class ConfigWithLiteral: """Precision level""" +class TestDSv4HybridRatioNormalization: + pattern = "WEC-|H-C-/W-" + compact = [0, 4, 128, 4, 0] + full = [0, 0, 4, 0, 128, 0, 4, 0, 0, 0] + + def test_derives_compact_and_full_ratios(self): + args = Namespace( + experimental_attention_variant="dsv4_hybrid", csa_compress_ratios=None + ) + config_kwargs = {} + + _normalize_dsv4_hybrid_csa_compress_ratios(args, config_kwargs, self.pattern) + + assert args.csa_compress_ratios == self.compact + assert config_kwargs["csa_compress_ratios"] == self.full + + def test_expands_compact_cli_ratios(self): + args = Namespace( + experimental_attention_variant="dsv4_hybrid", + csa_compress_ratios=list(self.compact), + ) + config_kwargs = {} + + _normalize_dsv4_hybrid_csa_compress_ratios(args, config_kwargs, self.pattern) + + assert args.csa_compress_ratios == self.compact + assert config_kwargs["csa_compress_ratios"] == self.full + + def test_rejects_ratio_that_disagrees_with_symbol(self): + args = Namespace( + experimental_attention_variant="dsv4_hybrid", + csa_compress_ratios=[0, 128, 128, 4, 0], + ) + + with pytest.raises(AssertionError, match="expected 4"): + _normalize_dsv4_hybrid_csa_compress_ratios(args, {}, self.pattern) + + def test_preserves_array_driven_d_ratio(self): + args = Namespace( + experimental_attention_variant="dsv4_hybrid", + csa_compress_ratios=[128, 4, 0], + ) + config_kwargs = {} + + _normalize_dsv4_hybrid_csa_compress_ratios(args, config_kwargs, "D-C/D") + + assert args.csa_compress_ratios == [128, 4, 0] + assert config_kwargs["csa_compress_ratios"] == [128, 0, 4, 0] + + class TestArgumentGroupFactoryBasic: """Test basic functionality of ArgumentGroupFactory.""" diff --git a/tests/unit_tests/test_fp8_utils.py b/tests/unit_tests/test_fp8_utils.py index 5be17f03c9f..fa7f8ff0760 100644 --- a/tests/unit_tests/test_fp8_utils.py +++ b/tests/unit_tests/test_fp8_utils.py @@ -10,6 +10,35 @@ from tests.unit_tests.test_utilities import Utils +@pytest.mark.skipif(not fp8_utils.HAVE_TE, reason="Transformer Engine is not installed") +@pytest.mark.parametrize( + ("is_init", "config_values", "te_helper"), + [ + ( + False, + {"fp8": "hybrid", "fp4": None, "fp8_param": False, "fp4_param": False}, + "fp8_autocast", + ), + ( + True, + {"fp8": None, "fp4": None, "fp8_param": True, "fp4_param": False}, + "fp8_model_init", + ), + ], +) +def test_get_fp8_disabled_context_uses_disabled_te_context(is_init, config_values, te_helper): + config = Mock(**config_values) + disabled_context = Mock() + + with patch.object( + fp8_utils.transformer_engine.pytorch, te_helper, return_value=disabled_context + ) as te_context: + result = fp8_utils.get_fp8_disabled_context(config, is_init=is_init) + + assert result is disabled_context + te_context.assert_called_once_with(enabled=False) + + class MockTELinear(nn.Module): """Mock TE Linear module for testing.""" diff --git a/tests/unit_tests/test_num_floating_point_operations.py b/tests/unit_tests/test_num_floating_point_operations.py index f358ada8fc1..66b0d714bc2 100644 --- a/tests/unit_tests/test_num_floating_point_operations.py +++ b/tests/unit_tests/test_num_floating_point_operations.py @@ -13,6 +13,7 @@ """ from types import SimpleNamespace +from unittest import mock import pytest import torch @@ -21,6 +22,7 @@ from megatron.training.training import ( consume_seqlen_stats_in_iteration, num_floating_point_operations, + set_seqlen_stats_in_iteration, update_seqlen_stats_from_cu_seqlens, ) @@ -29,6 +31,7 @@ def _reset_seqlen_accumulator(): """Tear down the per-iteration accumulator between tests.""" training_module._seqlen_stats_in_iteration = None training_module._seqlen_stats_active = False + training_module._seqlen_stats_are_global = False def _make_gpt_args( @@ -397,6 +400,26 @@ def test_update_accumulates_across_microbatches(self): assert total_real_tokens == 200 + 250 assert seqlen_squared_sum == 20000 + 42500 + def test_scheduler_global_stats_skip_all_reduce_and_reset(self): + set_seqlen_stats_in_iteration(450, 62500) + + assert training_module._seqlen_stats_active is True + assert training_module._seqlen_stats_are_global is True + with ( + mock.patch.object(torch.distributed, "is_initialized", return_value=True), + mock.patch.object( + training_module.mpu, "model_parallel_is_initialized", return_value=True + ), + mock.patch.object(torch.distributed, "all_reduce") as all_reduce, + ): + result = consume_seqlen_stats_in_iteration() + + assert result == (450, 62500) + all_reduce.assert_not_called() + assert training_module._seqlen_stats_active is False + assert training_module._seqlen_stats_are_global is False + assert training_module._seqlen_stats_in_iteration.tolist() == [0.0, 0.0] + def test_consume_resets_accumulator(self): cu = torch.tensor([0, 100, 200], dtype=torch.int32) update_seqlen_stats_from_cu_seqlens(cu) @@ -644,3 +667,185 @@ def test_dedup_across_topology(self, tp, cp, pp): f"topology tp={tp} cp={cp} pp={pp} dp={dp_size}: " f"got seqlen_squared_sum={seqlen_squared_sum}, expected {expected_sum_sq}" ) + + +def _make_dsv4_args(): + """Build a minimal DSv4 HybridModel, including one unified MTP layer.""" + return SimpleNamespace( + num_layers=4, + hidden_size=512, + num_attention_heads=8, + seq_length=256, + padded_vocab_size=1024, + swiglu=True, + ffn_hidden_size=2048, + kv_channels=64, + group_query_attention=False, + num_query_groups=8, + multi_latent_attention=True, + moe_router_topk=0, + moe_ffn_hidden_size=None, + moe_latent_size=None, + moe_shared_expert_intermediate_size=None, + mtp_num_layers=1, + experimental_attention_variant="dsv4_hybrid", + linear_key_head_dim=None, + linear_value_head_dim=None, + linear_num_key_heads=None, + linear_num_value_heads=None, + linear_conv_kernel_dim=None, + mamba_state_dim=128, + mamba_head_dim=64, + mamba_num_groups=8, + mamba_num_heads=128, + q_lora_rank=128, + qk_head_dim=32, + qk_pos_emb_head_dim=32, + kv_lora_rank=64, + v_head_dim=64, + o_lora_rank=64, + o_groups=2, + csa_window_size=64, + csa_compress_ratios=[0, 4, 128, 128, 0], + dsa_indexer_n_heads=4, + dsa_indexer_head_dim=32, + dsa_indexer_topk=16, + hybrid_layer_pattern="W-C-H-H-/W-", + ) + + +def _dsv4_golden_flops(args, total_tokens, seqlen_squared_sum): + """Independently calculate DSv4 HybridModel FLOPs from the pattern.""" + pattern = args.hybrid_layer_pattern.replace("|", "").replace("/", "") + n_r0 = pattern.count("W") + n_r4 = pattern.count("C") + n_r128 = pattern.count("H") + n_attention = n_r0 + n_r4 + n_r128 + n_mlp = pattern.count("-") + + q_projection = args.q_lora_rank * ( + args.hidden_size + args.num_attention_heads * args.v_head_dim + 1 + ) + kv_projection = args.hidden_size * args.v_head_dim + args.v_head_dim + output_projection = ( + args.num_attention_heads * args.v_head_dim * args.o_lora_rank + + args.o_groups * args.o_lora_rank * args.hidden_size + ) + attention_token_term = n_attention * ( + q_projection + kv_projection + output_projection + ) + + attention_token_term += ( + n_r0 * args.num_attention_heads * args.csa_window_size * args.v_head_dim * 2 + + n_r128 + * args.num_attention_heads + * args.csa_window_size + * args.v_head_dim + * 2 + ) + attention_core_term = n_r128 * args.num_attention_heads * args.v_head_dim / 128 + attention_token_term += ( + n_r4 * args.hidden_size * (2 * args.v_head_dim) * 2 + + n_r128 * args.hidden_size * args.v_head_dim * 2 + ) + + if n_r4: + effective_topk = min(args.dsa_indexer_topk, args.seq_length // 4) + average_compressed_tokens = effective_topk * ( + 1 - effective_topk * 4 / (2 * args.seq_length) + ) + attention_token_term += ( + n_r4 + * args.num_attention_heads + * (args.csa_window_size + average_compressed_tokens) + * args.v_head_dim + * 2 + + n_r4 * args.hidden_size * (2 * args.dsa_indexer_head_dim) * 2 + + n_r4 + * args.q_lora_rank + * args.dsa_indexer_n_heads + * args.dsa_indexer_head_dim + + n_r4 * args.hidden_size * args.dsa_indexer_n_heads + ) + attention_core_term += ( + n_r4 * args.dsa_indexer_n_heads * args.dsa_indexer_head_dim / 4 + ) + + mlp_expansion = args.ffn_hidden_size / args.hidden_size + swiglu_scale = 3 / 2 if args.swiglu else 1 + mlp_forward = ( + n_mlp + * 4 + * mlp_expansion + * swiglu_scale + * total_tokens + * args.hidden_size**2 + ) + mtp_forward = ( + 2 + * args.mtp_num_layers + * (3 * args.hidden_size + 2 * args.hidden_size**2) + * total_tokens + ) + logits_forward = ( + 2 + * total_tokens + * args.hidden_size + * args.padded_vocab_size + * (1 + args.mtp_num_layers) + ) + attention_forward = 2 * ( + attention_token_term * total_tokens + + attention_core_term * seqlen_squared_sum + ) + return 3 * (attention_forward + mlp_forward + mtp_forward + logits_forward) + + +class TestDSv4Hybrid: + """DSv4 HybridModel FLOPs against an independent pattern-based golden.""" + + def test_bshd(self): + args = _make_dsv4_args() + batch_size = 2 + total_tokens = batch_size * args.seq_length + seqlen_squared_sum = batch_size * args.seq_length**2 + + actual = num_floating_point_operations(args, batch_size) + expected = _dsv4_golden_flops(args, total_tokens, seqlen_squared_sum) + + assert actual == pytest.approx(expected) + + def test_thd(self): + args = _make_dsv4_args() + batch_size = 2 + packed_lengths = [64, 64, 128, 256] + total_tokens = sum(packed_lengths) + seqlen_squared_sum = sum(length**2 for length in packed_lengths) + + actual = num_floating_point_operations( + args, + batch_size, + total_real_tokens_in_batch=total_tokens, + seqlen_squared_sum_in_batch=seqlen_squared_sum, + ) + expected = _dsv4_golden_flops(args, total_tokens, seqlen_squared_sum) + + assert actual == pytest.approx(expected) + assert actual < num_floating_point_operations(args, batch_size) + + def test_mtp_projection_norm_accounting(self): + args = _make_dsv4_args() + without_mtp_scaffolding = SimpleNamespace(**vars(args)) + without_mtp_scaffolding.mtp_num_layers = 0 + batch_size = 2 + total_tokens = batch_size * args.seq_length + + delta = num_floating_point_operations( + args, batch_size + ) - num_floating_point_operations(without_mtp_scaffolding, batch_size) + expected = 3 * ( + 2 * total_tokens * (3 * args.hidden_size + 2 * args.hidden_size**2) + + 2 * total_tokens * args.hidden_size * args.padded_vocab_size + ) + + assert delta == expected diff --git a/tests/unit_tests/test_sequence_packing.py b/tests/unit_tests/test_sequence_packing.py new file mode 100644 index 00000000000..85faa14fc50 --- /dev/null +++ b/tests/unit_tests/test_sequence_packing.py @@ -0,0 +1,820 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import random +import sys +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from megatron.core import parallel_state +from megatron.core.datasets.data_schedule import ( + DpBalancedScheduler, + _build_thd_padding_mask, + _get_scheduler_max_real_num_seqs, + _sanitize_thd_padding_values, + get_batch_on_this_rank_for_sequence_packing, + wrap_data_iterator, +) +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.rerun_state_machine import RerunDataIterator +from megatron.training.global_vars import unset_global_variables +from tests.unit_tests.test_utilities import Utils + + +def _scheduler_pg_collection(): + """Build the process groups consumed by the packing scheduler.""" + return ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['tp', 'pp', 'cp', 'dp', 'dp_cp'] + ) + + +def test_te_thd_partition_helper_delegates_to_transformer_engine(monkeypatch): + from megatron.core.extensions import transformer_engine as te_extension + + calls = [] + expected = torch.tensor([3, 1], dtype=torch.int64) + + def _partition(cu_seqlens, total_tokens, cp_size, cp_rank): + calls.append((cu_seqlens, total_tokens, cp_size, cp_rank)) + return expected + + monkeypatch.setattr(te_extension, "is_te_min_version", lambda _version: True) + monkeypatch.setitem( + sys.modules, + "transformer_engine_torch", + SimpleNamespace(thd_get_partitioned_indices=_partition), + ) + cu_seqlens = torch.tensor([0, 2, 6], dtype=torch.int32) + + actual = te_extension.get_thd_partitioned_indices(cu_seqlens, 6, 2, 1) + + assert actual is expected + assert calls == [(cu_seqlens, 6, 2, 1)] + + +def test_scheduler_thd_padding_mask_from_cu_seqlens(): + cu_seqlens = torch.tensor([0, 3, 5], dtype=torch.int32) + cu_seqlens_padded = torch.tensor([0, 4, 8], dtype=torch.int32) + + padding_mask = _build_thd_padding_mask(cu_seqlens, cu_seqlens_padded) + + assert torch.equal( + padding_mask, torch.tensor([False, False, False, True, False, False, True, True]) + ) + + +def test_scheduler_sanitizes_thd_padding_values(): + padding_mask = torch.tensor([False, False, True, False, True]) + batch = { + 'tokens': torch.tensor([11, 12, -1, 21, -1], dtype=torch.int64), + 'labels': torch.tensor([12, 13, -1, 22, -1], dtype=torch.int64), + 'loss_mask': torch.ones(5, dtype=torch.float32), + 'position_ids': torch.tensor([0, 1, 2, 0, 1], dtype=torch.int64), + } + + _sanitize_thd_padding_values(batch, padding_mask) + + assert torch.equal(batch['tokens'], torch.tensor([11, 12, 0, 21, 0])) + assert torch.equal(batch['labels'], torch.tensor([12, 13, 0, 22, 0])) + assert torch.equal(batch['loss_mask'], torch.tensor([1.0, 1.0, 0.0, 1.0, 0.0])) + assert torch.equal(batch['position_ids'], torch.tensor([0, 1, 0, 0, 0])) + + +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 + + +@pytest.mark.parametrize( + "rank, expected_tokens, expected_padding_mask", + [ + (2, torch.arange(8, 12, dtype=torch.int64), torch.zeros(4, dtype=torch.bool)), + (3, torch.zeros(4, dtype=torch.int64), torch.ones(4, dtype=torch.bool)), + ], +) +def test_dsv4_thd_cp_slice_uses_static_partition_total( + rank, expected_tokens, expected_padding_mask +): + 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), + "padding_mask": torch.zeros(12, dtype=torch.bool), + "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=rank), + keys=("tokens", "position_ids", "labels", "loss_mask", "padding_mask"), + cp_partition_mode="contiguous", + partition_total_tokens=16, + ) + + assert torch.equal(batch["tokens"], expected_tokens) + assert torch.equal(batch["position_ids"], expected_tokens) + expected_labels = expected_tokens + 100 if rank == 2 else expected_tokens + assert torch.equal(batch["labels"], expected_labels) + expected_loss_mask = ( + torch.ones(4, dtype=torch.float32) + if rank == 2 + else torch.zeros(4, dtype=torch.float32) + ) + assert torch.equal(batch["loss_mask"], expected_loss_mask) + assert torch.equal(batch["padding_mask"], expected_padding_mask) + 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)) + + +def test_dsv4_thd_contiguous_cp_slice_rejects_invalid_static_totals(): + from megatron.core.datasets.data_schedule_utils import get_cp_slice_for_thd + + batch = { + "tokens": torch.arange(12, dtype=torch.int64), + "cu_seqlens_padded": torch.tensor([0, 12], dtype=torch.int32), + } + with pytest.raises(RuntimeError, match="smaller than tokens length"): + get_cp_slice_for_thd( + batch, + _MockCPGroup(size=4, rank=0), + keys=("tokens",), + cp_partition_mode="contiguous", + partition_total_tokens=8, + ) + with pytest.raises(RuntimeError, match="divisible"): + get_cp_slice_for_thd( + batch, + _MockCPGroup(size=5, rank=0), + keys=("tokens",), + cp_partition_mode="contiguous", + ) + + +def test_packed_batch_preserves_original_and_padded_cu_seqlens(): + Utils.initialize_model_parallel(1, 1) + + try: + device = torch.device("cuda", torch.cuda.current_device()) + tokens = torch.arange(8, dtype=torch.int64, device=device) + batch = { + 'tokens': tokens, + 'labels': tokens + 1, + 'loss_mask': torch.ones(8, dtype=torch.float32, device=device), + 'position_ids': torch.arange(8, dtype=torch.int64, device=device), + 'cu_seqlens': torch.tensor([0, 3, 5], dtype=torch.int32, device=device), + 'cu_seqlens_padded': torch.tensor([0, 4, 8], dtype=torch.int32, device=device), + 'max_seqlen': torch.tensor([4], dtype=torch.int32, device=device), + } + + *_, packed_seq_params, padding_mask = get_batch_on_this_rank_for_sequence_packing( + iter([batch]), pg_collection=_scheduler_pg_collection() + ) + + torch.testing.assert_close(packed_seq_params.cu_seqlens_q, batch['cu_seqlens']) + torch.testing.assert_close( + packed_seq_params.cu_seqlens_q_padded, batch['cu_seqlens_padded'] + ) + assert torch.equal( + padding_mask, + torch.tensor([[False, False, False, True, False, False, True, True]], device=device), + ) + finally: + Utils.destroy_model_parallel() + + +def test_dp_balanced_scheduler_can_split_group_zero(): + scheduler = DpBalancedScheduler( + max_seqlen_per_dp_cp_rank=8, cp_size=1, dp_size=2, microbatch_group_size_per_vp_stage=None + ) + + assert scheduler.get_groups_and_subsamples([(0, 2), (1, 2)]) == [[[0], [1]]] + + +def test_dp_balanced_scheduler_reserves_dummy_tail_capacity(): + config = SimpleNamespace( + thd_max_packed_sequences=3, + pad_packed_seq_alignment="max", + pad_packed_seq_by_appending_dummy_seq=True, + ) + assert _get_scheduler_max_real_num_seqs(config) == 2 + + scheduler = DpBalancedScheduler( + max_seqlen_per_dp_cp_rank=16, + cp_size=1, + dp_size=1, + microbatch_group_size_per_vp_stage=None, + max_num_seqs=2, + ) + assert scheduler.get_groups_and_subsamples( + [(0, 1), (1, 1), (2, 1), (3, 1)] + ) == [[[0, 1]], [[2, 3]]] + + +def test_get_batch_applies_static_thd_padding_from_config(): + Utils.initialize_model_parallel(1, 1) + try: + device = torch.device("cuda", torch.cuda.current_device()) + tokens = torch.arange(8, dtype=torch.int64, device=device) + batch = { + 'tokens': tokens, + 'labels': tokens + 1, + 'loss_mask': torch.ones(8, dtype=torch.float32, device=device), + 'position_ids': torch.arange(8, dtype=torch.int64, device=device), + 'cu_seqlens': torch.tensor([0, 3, 5], dtype=torch.int32, device=device), + 'cu_seqlens_padded': torch.tensor([0, 4, 8], dtype=torch.int32, device=device), + 'max_seqlen': torch.tensor([4], dtype=torch.int32, device=device), + } + config = SimpleNamespace( + pad_packed_seq_alignment="max", + max_seqlen_per_dp_cp_rank=16, + thd_max_packed_sequences=4, + cuda_graph_impl="transformer_engine", + pad_packed_seq_by_appending_dummy_seq=True, + ) + + tokens, *_, packed_seq_params, padding_mask = ( + get_batch_on_this_rank_for_sequence_packing( + iter([batch]), + pg_collection=_scheduler_pg_collection(), + config=config, + ) + ) + + assert tokens.shape == (1, 16) + assert packed_seq_params.cu_seqlens_q.shape == (5,) + assert packed_seq_params.cu_seqlens_q.tolist() == [0, 3, 5, 16, 16] + assert packed_seq_params.cu_seqlens_q_padded.tolist() == [0, 4, 8, 16, 16] + assert padding_mask.shape == (1, 16) + assert padding_mask[0, 8:].all() + finally: + Utils.destroy_model_parallel() + + +class MockVariableLengthSequencePackingDataIterator: + """ + Mock data iterator for testing get_batch_on_this_rank_for_sequence_packing. + + Generates variable-length (THD format) packed sequences with deterministic + data for verification across parallel ranks. + """ + + def __init__( + self, + total_seq_length: int, + sequence_lengths: list, + local_cp_size: int = None, + device: str = "cuda", + seed: int = 42, + ): + """ + Args: + total_seq_length: Total length of packed sequences + sequence_lengths: List of individual sequence lengths (variable-length). + If None, generates random variable lengths. + device: Device to create tensors on + seed: Random seed for reproducibility + """ + self.total_seq_length = total_seq_length + self.sequence_lengths = sequence_lengths + self.local_cp_size = local_cp_size + self.device = device + self.seed = seed + assert ( + sum(self.sequence_lengths) == total_seq_length + ), f"Sequence lengths sum {sum(self.sequence_lengths)} != total {total_seq_length}" + + def __iter__(self): + """Interface for the data iterator.""" + return self + + def __next__(self): + """Generate a mock batch with variable-length THD format.""" + dev = self.device + torch.manual_seed(self.seed) + torch.cuda.manual_seed(self.seed) + + tokens = torch.randint(0, 16384, (self.total_seq_length,), dtype=torch.int64, device=dev) + + # Create position_ids that reset for each sequence (THD format) + position_ids = [] + for seq_len in self.sequence_lengths: + position_ids.extend(range(seq_len)) + position_ids = torch.tensor(position_ids, dtype=torch.int64, device=dev) + + # Labels are tokens shifted by 1 for easy verification + labels = tokens + 1 + + # Loss mask: 1.0 for all positions except padding (none here) + loss_mask = torch.ones(self.total_seq_length, dtype=torch.float32, device=dev) + + # Create cu_seqlens for variable-length packed sequences + cu_seqlens = [0] + for seq_len in self.sequence_lengths: + cu_seqlens.append(cu_seqlens[-1] + seq_len) + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=dev) + cu_seqlens_padded = cu_seqlens.clone() + + max_seqlen = torch.tensor([max(self.sequence_lengths)], dtype=torch.int32, device=dev) + + batch = { + "tokens": tokens, + "position_ids": position_ids, + "labels": labels, + "loss_mask": loss_mask, + "cu_seqlens": cu_seqlens, + "cu_seqlens_padded": cu_seqlens_padded, + "max_seqlen": max_seqlen, + } + + if not ( + parallel_state.is_pipeline_first_stage(ignore_virtual=True) + or parallel_state.is_pipeline_last_stage(ignore_virtual=True) + ): + batch["tokens"] = None + batch["position_ids"] = None + batch["labels"] = None + batch["loss_mask"] = None + + if self.local_cp_size is not None: + batch["local_cp_size"] = torch.tensor( + [self.local_cp_size], dtype=torch.int32, device=dev + ) + + return batch + + +def _gather_tensor_from_tp_group(tensor): + """Gather tensors from all TP ranks for comparison.""" + assert tensor is not None, "Tensor should not be None" + tp_size = parallel_state.get_tensor_model_parallel_world_size() + gathered = [torch.zeros_like(tensor) for _ in range(tp_size)] + torch.distributed.all_gather( + gathered, tensor, group=parallel_state.get_tensor_model_parallel_group() + ) + return gathered + + +def _gather_tensor_from_all_ranks(tensor): + """Gather tensors from all PP ranks for comparison.""" + assert tensor is not None, "Tensor should not be None" + if type(tensor) is int: + tensor = torch.tensor(tensor, dtype=torch.int32, device=torch.cuda.current_device()) + gathered = [torch.zeros_like(tensor) for _ in range(torch.distributed.get_world_size())] + torch.distributed.all_gather(gathered, tensor) + return gathered + + +@pytest.mark.parametrize( + ("tp", "pp", "cp"), + [ + (1, 1, 1), # Basic case: no parallelism + (2, 1, 1), # Tensor parallel only + (1, 2, 1), # Pipeline parallel only + (2, 2, 1), # TP + PP + (1, 1, 2), # CP only + (2, 1, 2), # TP + CP + (1, 2, 2), # PP + CP + (1, 4, 1), # Has middle pp stage + ], +) +def test_get_batch_on_this_rank_for_sequence_packing(tp, pp, cp): + """ + Test get_batch_on_this_rank_for_sequence_packing function with variable-length THD format. + + This test verifies: + 1. TP ranks: All ranks within a TP group receive identical data after broadcast + 2. PP ranks: Middle PP ranks have the same packed_seq_params as first/last stages + 3. CP ranks: Data is correctly partitioned with proper shape and values + 4. Variable-length (THD) format: Different sequence lengths are handled correctly + """ + args = SimpleNamespace() + args.tensor_model_parallel_size = tp + args.pipeline_model_parallel_size = pp + args.context_parallel_size = cp + args.virtual_pipeline_model_parallel_size = None + args.data_parallel_size = 8 // (tp * pp * cp) + args.seq_length = 8192 + + # Skip invalid configurations + if args.data_parallel_size < 1: + raise ValueError(f"Invalid config: tp={tp}, pp={pp}, cp={cp} exceeds world size 8") + + # Initialize model parallel + Utils.initialize_model_parallel(tp, pp, None, context_parallel_size=cp) + + try: + # Create mock data iterator with variable-length sequences + # Only TP rank 0 needs the iterator; other TP ranks pass None + tp_rank = parallel_state.get_tensor_model_parallel_rank() + if tp_rank == 0: + # Use deterministic seed based on DP rank so same data within TP/PP/CP group + dp_rank = parallel_state.get_data_parallel_rank() + sequence_lengths = [1024, 2048, 512, 1536, 3072] + assert ( + sum(sequence_lengths) == args.seq_length + ), f"Sequence lengths sum {sum(sequence_lengths)} != total {args.seq_length}" + data_iterator = iter( + MockVariableLengthSequencePackingDataIterator( + total_seq_length=args.seq_length, + sequence_lengths=sequence_lengths, # Variable lengths, sum=8192 + seed=42 + dp_rank, # Same seed within PP/CP group + ) + ) + else: + # Non-TP-rank-0 ranks don't need the iterator + data_iterator = None + + # Call the function under test + result = get_batch_on_this_rank_for_sequence_packing( + data_iterator=data_iterator, + pg_collection=_scheduler_pg_collection(), + mtp_on_this_rank=False, + vp_stage=None, + ) + + # Unpack the result. Scheduler THD always returns padding_mask. + tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params, padding_mask = ( + result + ) + + # Get parallel state info + tp_rank = parallel_state.get_tensor_model_parallel_rank() + pp_rank = parallel_state.get_pipeline_model_parallel_rank() + cp_rank = parallel_state.get_context_parallel_rank() + is_first_stage = parallel_state.is_pipeline_first_stage(ignore_virtual=True) + is_last_stage = parallel_state.is_pipeline_last_stage(ignore_virtual=True) + is_first_or_last = is_first_stage or is_last_stage + + assert padding_mask is not None + assert padding_mask.dtype == torch.bool + assert padding_mask.dim() == 2 + assert padding_mask.size(0) == 1 + assert not padding_mask.any(), "Mock data has no per-sequence padding." + + # ===================================================================== + # TEST 1: Verify data based on pipeline stage + # ===================================================================== + if is_first_stage: + assert tokens is not None, "First stage should have tokens" + assert position_ids is not None, "First stage should have position_ids" + assert tokens.dim() == 2, "Tokens should be 2D (batch, seq)" + assert position_ids.dim() == 2, "Position IDs should be 2D (batch, seq)" + assert tokens.size(0) == 1, "batch should be 1 in THD format" + assert position_ids.size(0) == 1, "batch should be 1 in THD format" + else: + assert tokens is None, "Non-first stage should not have tokens" + assert position_ids is None, "Non-first stage should not have position_ids" + + if is_last_stage: + assert labels is not None, "Last stage should have labels" + assert loss_mask is not None, "Last stage should have loss_mask" + assert labels.dim() == 2, "Labels should be 2D (batch, seq)" + assert loss_mask.dim() == 2, "Loss mask should be 2D (batch, seq)" + assert labels.size(0) == 1, "batch should be 1 in THD format" + assert loss_mask.size(0) == 1, "batch should be 1 in THD format" + else: + assert labels is None, "Non-last stage should not have labels" + assert loss_mask is None, "Non-last stage should not have loss_mask" + + # ===================================================================== + # TEST 2: Verify all ranks have consistent packed_seq_params + # ===================================================================== + assert packed_seq_params is not None + assert packed_seq_params.qkv_format == "thd" + + test_keys = [ + "cu_seqlens_q", + "cu_seqlens_q_padded", + "max_seqlen_q", + "cu_seqlens_kv", + "cu_seqlens_kv_padded", + "max_seqlen_kv", + ] + for key in test_keys: + tensor = getattr(packed_seq_params, key) + assert tensor is not None + gathered_tensor = _gather_tensor_from_all_ranks(tensor) + for i in range(1, len(gathered_tensor)): + assert torch.equal( + gathered_tensor[0], gathered_tensor[i] + ), f"Rank 0 and rank {i} have different {key}" + + # ===================================================================== + # TEST 3: Verify TP ranks receive identical data after broadcast + # ===================================================================== + if tp > 1: + test_tensors = [padding_mask] + if is_first_stage: + test_tensors.extend([tokens, position_ids]) + if is_last_stage: + test_tensors.extend([labels, loss_mask]) + + for tensor in test_tensors: + gathered_tensors = _gather_tensor_from_tp_group(tensor) + for i in range(1, tp): + assert torch.equal( + gathered_tensors[0], gathered_tensors[i] + ), f"TP rank 0 and rank {i} have different data" + + # ===================================================================== + # TEST 4: Verify CP partitioning + # ===================================================================== + if cp > 1: + # With CP, the sequence should be partitioned + expected_seq_len = args.seq_length // cp + + if is_first_stage: + actual_seq_len = tokens.shape[1] + assert ( + actual_seq_len == expected_seq_len + ), f"CP partitioned tokens have wrong shape: {actual_seq_len} != {expected_seq_len}" + + # Verify labels only if all CP ranks are at last stage + if is_last_stage: + actual_seq_len = labels.shape[1] + assert ( + actual_seq_len == expected_seq_len + ), f"CP partitioned labels have wrong shape: {actual_seq_len} != {expected_seq_len}" + + actual_seq_len = padding_mask.shape[1] + assert ( + actual_seq_len == expected_seq_len + ), f"CP partitioned padding_mask has wrong shape: {actual_seq_len} != {expected_seq_len}" + + finally: + Utils.destroy_model_parallel() + unset_global_variables() + + +@pytest.mark.parametrize( + ("tp", "pp", "cp", "vpp", "scheduler_type", "mtp_vpp"), + [ + (1, 1, 8, None, "dp_balanced", False), + (2, 1, 4, None, "dp_balanced", False), + (2, 4, 1, None, "dp_balanced", False), + (2, 2, 1, None, "dp_balanced", False), + (1, 4, 1, 4, "dp_balanced", False), + (1, 4, 1, 4, "dp_balanced", True), + ], +) +def test_wrap_dataloader(tp, pp, cp, vpp, scheduler_type, mtp_vpp, monkeypatch): + ''' + Test wrap_dataloader function with different scheduler types. + ''' + args = SimpleNamespace() + args.tensor_model_parallel_size = tp + args.pipeline_model_parallel_size = pp + args.context_parallel_size = cp + args.virtual_pipeline_model_parallel_size = None + args.data_parallel_size = 8 // (tp * pp * cp) + args.seq_length = 8192 + args.max_seqlen_per_dp_cp_rank = 8192 + + # Skip invalid configurations + if args.data_parallel_size < 1: + raise ValueError(f"Invalid config: tp={tp}, pp={pp}, cp={cp} exceeds world size 8") + + def _create_single_sample(seq_len): + # hard code the padding size to 16 + pad_size = 16 + seq_len_padded = ((seq_len + pad_size - 1) // pad_size) * pad_size + device = torch.device("cuda", torch.cuda.current_device()) + tokens = torch.randint(0, 128, (seq_len_padded,), dtype=torch.int64, device=device) + labels = tokens + 1 + position_ids = torch.arange(seq_len_padded, dtype=torch.int64, device=device) + loss_mask = torch.ones(seq_len_padded, dtype=torch.float32, device=device) + loss_mask[0:seq_len] = 1 + loss_mask[seq_len:] = 0 + cu_seqlens = torch.tensor([0, seq_len_padded], dtype=torch.int32, device=device) + + return { + 'tokens': tokens, + 'labels': labels, + 'loss_mask': loss_mask, + 'position_ids': position_ids, + 'cu_seqlens': cu_seqlens, + } + + # Initialize model parallel + Utils.initialize_model_parallel(tp, pp, vpp, context_parallel_size=cp) + + global_batch_size = 64 + micro_batch_size = 1 + nums = [random.randint(2048, args.seq_length) for _ in range(global_batch_size)] # 64 sequences + + config = SimpleNamespace() + config.max_seqlen_per_dp_cp_rank = args.max_seqlen_per_dp_cp_rank + config.microbatch_group_size_per_vp_stage = pp + config.virtual_pipeline_model_parallel_size = vpp + config.sequence_packing_scheduler = scheduler_type + config.pipeline_model_parallel_layout = object() if mtp_vpp else None + config.mtp_num_layers = 1 if mtp_vpp else None + + dp_rank = parallel_state.get_data_parallel_rank() + dp_size = parallel_state.get_data_parallel_world_size() + + pp_rank = parallel_state.get_pipeline_model_parallel_rank() + tp_rank = parallel_state.get_tensor_model_parallel_rank() + + is_pp_first = pp_rank == 0 + is_pp_last = pp_rank == pp - 1 + is_tp_first = tp_rank == 0 + + if mtp_vpp: + mtp_pp_rank = 1 + mtp_vp_stage = 2 + + def _mock_mtp_on_this_rank(*, ignore_virtual, vp_stage=None, **_kwargs): + return pp_rank == mtp_pp_rank and (ignore_virtual or vp_stage == mtp_vp_stage) + + monkeypatch.setattr( + "megatron.core.datasets.data_schedule.mtp_is_on_rank", _mock_mtp_on_this_rank + ) + + num_micro_batches_old = global_batch_size // micro_batch_size // dp_size + + if is_tp_first: + samples = [ + _create_single_sample(num) + for num in nums[dp_rank * num_micro_batches_old : (dp_rank + 1) * num_micro_batches_old] + ] + data_iterator = RerunDataIterator(iter(samples)) + else: + data_iterator = None + + if is_tp_first: + if vpp is not None and vpp > 1: + data_iterator = [data_iterator] + [None for _ in range(vpp - 1)] + try: + # Call the function under test + ( + new_data_iterator, + num_micro_batches, + num_total_tokens_this_global_batch, + sequence_square_sum_this_global_batch, + ) = wrap_data_iterator( + data_iterator, + config, + num_micro_batches_old, + _scheduler_pg_collection(), + ) + + # check the result + assert type(num_micro_batches) is int + assert ( + type(num_total_tokens_this_global_batch) is float + or type(num_total_tokens_this_global_batch) is np.float32 + ) + assert ( + type(sequence_square_sum_this_global_batch) is float + or type(sequence_square_sum_this_global_batch) is np.float32 + ) + + def _check_batch(batch_all, batch_keys): + for batch in batch_all: + assert set(batch) == set( + batch_keys + ), f"batch keys: {set(batch)} != expected keys: {set(batch_keys)}" + for key in batch_keys: + assert batch[key] is not None + + if is_tp_first: + metadata_keys = ["cu_seqlens", "max_seqlen", "cu_seqlens_padded"] + + def _expected_keys(vp_stage=None): + keys = list(metadata_keys) + is_mtp_stage = mtp_vpp and pp_rank == mtp_pp_rank and vp_stage == mtp_vp_stage + if (is_pp_first and (vp_stage is None or vp_stage == 0)) or is_mtp_stage: + keys.extend(["tokens", "position_ids"]) + if (is_pp_last and (vp_stage is None or vp_stage == vpp - 1)) or is_mtp_stage: + keys.extend(["labels", "loss_mask"]) + return keys + + token_batches = None + if vpp is not None and vpp > 1: + # check metadata for all stages (save batches to avoid re-consuming iterators) + all_stage_batches = [] + for vp_stage, temp_data_iterator in enumerate(new_data_iterator): + stage_batch = [next(temp_data_iterator) for _ in range(num_micro_batches)] + all_stage_batches.append(stage_batch) + _check_batch(stage_batch, _expected_keys(vp_stage)) + if is_pp_first: + token_batches = all_stage_batches[0] + else: + # non-VPP: single iterator + batch_all = [next(new_data_iterator) for _ in range(num_micro_batches)] + _check_batch(batch_all, _expected_keys()) + if is_pp_first: + token_batches = batch_all + + # CHECK TOKEN SUM ON FIRST PP RANK + # Note: data_iterator is consumed by wrap_data_iterator, new_data_iterator is consumed above. + # Use `samples` for before-wrap and the first data-carrying iterator after wrapping. + if is_pp_first: + # Compute token sum before wrap + token_sum_before = torch.tensor(0, dtype=torch.int64, device='cuda') + for sample in samples: + token_sum_before += sample['tokens'].long().sum() + + # Compute token sum after wrap. + token_sum_after = torch.tensor(0, dtype=torch.int64, device='cuda') + for batch in token_batches: + token_sum_after += batch['tokens'].long().sum() + + # Reduce sum across dp_cp group and verify equality + dp_cp_group = parallel_state.get_data_parallel_group(with_context_parallel=False) + torch.distributed.all_reduce( + token_sum_before, op=torch.distributed.ReduceOp.SUM, group=dp_cp_group + ) + torch.distributed.all_reduce( + token_sum_after, op=torch.distributed.ReduceOp.SUM, group=dp_cp_group + ) + + assert ( + token_sum_before == token_sum_after + ), f"Token sum mismatch: before={token_sum_before.item()}, after={token_sum_after.item()}" + + else: + if vpp is not None and vpp > 1: + assert type(new_data_iterator) is list and len(new_data_iterator) == vpp + for data_iterator in new_data_iterator: + assert data_iterator is None + else: + assert new_data_iterator is None + + finally: + Utils.destroy_model_parallel() + unset_global_variables() + + +def test_wrapped_batch_with_pipeline_and_context_parallel(): + """Exercise scheduler field filtering followed by PP-aware CP slicing.""" + Utils.initialize_model_parallel(1, 2, None, context_parallel_size=2) + + try: + device = torch.device("cuda", torch.cuda.current_device()) + dp_rank = parallel_state.get_data_parallel_rank() + tokens = torch.arange(16, dtype=torch.int64, device=device) + dp_rank * 100 + sample = { + 'tokens': tokens, + 'labels': tokens + 1, + 'loss_mask': torch.ones(16, dtype=torch.float32, device=device), + 'position_ids': torch.arange(16, dtype=torch.int64, device=device), + 'original_seq_len': torch.tensor([12], dtype=torch.int32, device=device), + 'padded_seq_len': torch.tensor([16], dtype=torch.int32, device=device), + } + config = SimpleNamespace( + max_seqlen_per_dp_cp_rank=8, + microbatch_group_size_per_vp_stage=None, + virtual_pipeline_model_parallel_size=None, + sequence_packing_scheduler="dp_balanced", + pipeline_model_parallel_layout=None, + mtp_num_layers=None, + ) + + packed_iterator, num_microbatches, token_sum, squared_sum = wrap_data_iterator( + RerunDataIterator(iter([sample])), + config, + 1, + _scheduler_pg_collection(), + ) + assert num_microbatches == 1 + assert token_sum == 24.0 + assert squared_sum == 288.0 + + tokens, labels, loss_mask, _, position_ids, packed_seq_params, padding_mask = ( + get_batch_on_this_rank_for_sequence_packing( + packed_iterator, pg_collection=_scheduler_pg_collection() + ) + ) + is_first = parallel_state.is_pipeline_first_stage() + assert (tokens is not None) == is_first + assert (position_ids is not None) == is_first + assert (labels is None) == is_first + assert (loss_mask is None) == is_first + assert packed_seq_params.qkv_format == "thd" + torch.testing.assert_close( + packed_seq_params.cu_seqlens_q, torch.tensor([0, 12], dtype=torch.int32, device=device) + ) + torch.testing.assert_close( + packed_seq_params.cu_seqlens_q_padded, + torch.tensor([0, 16], dtype=torch.int32, device=device), + ) + assert padding_mask.shape == (1, 8) + finally: + Utils.destroy_model_parallel() + unset_global_variables() diff --git a/tests/unit_tests/training/models/test_hybrid_builder.py b/tests/unit_tests/training/models/test_hybrid_builder.py index 9984e224ce3..bc879620c03 100644 --- a/tests/unit_tests/training/models/test_hybrid_builder.py +++ b/tests/unit_tests/training/models/test_hybrid_builder.py @@ -206,6 +206,20 @@ def test_spec_already_module_spec_used_directly(self, mock_model, *_): call_kwargs = mock_model.call_args.kwargs assert call_kwargs["hybrid_stack_spec"] is module_spec + @patch("megatron.training.models.hybrid.calculate_padded_vocab_size") + @patch("megatron.training.models.hybrid.is_pp_last_stage", return_value=True) + @patch("megatron.training.models.hybrid.is_pp_first_stage", return_value=True) + @patch("megatron.training.models.hybrid.HybridModel") + def test_callable_spec_receives_transformer_config(self, mock_model, *_): + module_spec = ModuleSpec(module=object) + spec_factory = Mock(return_value=module_spec) + self.config.__dict__["hybrid_stack_spec"] = spec_factory + + self.builder.build_model(self.pg, pre_process=True, post_process=True) + + spec_factory.assert_called_once_with(self.config.transformer) + assert mock_model.call_args.kwargs["hybrid_stack_spec"] is module_spec + @patch("megatron.training.models.hybrid.calculate_padded_vocab_size") @patch("megatron.training.models.hybrid.is_pp_last_stage", return_value=True) @patch("megatron.training.models.hybrid.is_pp_first_stage", return_value=True) diff --git a/tests/unit_tests/training/test_train_step_schedule_plumbing.py b/tests/unit_tests/training/test_train_step_schedule_plumbing.py index 6dcb407f920..44a375031c3 100644 --- a/tests/unit_tests/training/test_train_step_schedule_plumbing.py +++ b/tests/unit_tests/training/test_train_step_schedule_plumbing.py @@ -5,6 +5,9 @@ from types import SimpleNamespace from unittest import mock +import torch + +from megatron.core.process_groups_config import ProcessGroupCollection from megatron.training import training as training_mod @@ -67,3 +70,300 @@ def test_train_step_forwards_schedule_plumbing(): def test_train_step_defaults_to_none(): captured = _run() assert captured["p2p_communicator"] is None and captured["pg_collection"] is None + + +def test_train_step_sets_scheduler_global_seqlen_stats_after_forward(): + args = SimpleNamespace( + save_params_interval=None, + save_activations_interval=None, + save_tokens_per_expert_interval=None, + save_wgrads_interval=None, + save_dgrads_interval=None, + reuse_grad_buf_for_mxfp8_param_ag=False, + overlap_param_gather=False, + seq_length=8, + micro_batch_size=1, + decoder_seq_length=None, + empty_unused_memory_level=0, + ) + pg_collection = ProcessGroupCollection() + model = [ + SimpleNamespace( + force_all_reduce=False, + zero_grad_buffer=lambda: None, + pg_collection=pg_collection, + ) + ] + config = SimpleNamespace(sequence_packing_scheduler=object()) + events = [] + original_setter = training_mod.set_seqlen_stats_in_iteration + + def forward_backward(**kwargs): + events.append("forward") + training_mod.update_seqlen_stats_from_cu_seqlens( + torch.tensor([0, 7], dtype=torch.int32) + ) + return [] + + def record_setter(total_real_tokens, seqlen_squared_sum): + events.append(("setter", total_real_tokens, seqlen_squared_sum)) + return original_setter(total_real_tokens, seqlen_squared_sum) + + training_mod._seqlen_stats_in_iteration = None + training_mod._seqlen_stats_active = False + training_mod._seqlen_stats_are_global = False + try: + with ( + mock.patch.object(training_mod, "get_args", return_value=args), + mock.patch.object(training_mod, "get_timers", return_value=mock.MagicMock()), + mock.patch.object(training_mod, "get_rerun_state_machine", return_value=_Rerun()), + mock.patch.object(training_mod, "get_num_microbatches", return_value=8), + mock.patch.object(training_mod, "has_nvidia_modelopt", False), + mock.patch.object( + training_mod, + "wrap_data_iterator", + return_value=(iter([]), 3, 450, 62500), + ), + mock.patch.object( + training_mod, "set_seqlen_stats_in_iteration", side_effect=record_setter + ), + ): + training_mod.train_step( + forward_step_func=lambda *args, **kwargs: None, + data_iterator=iter([]), + model=model, + optimizer=SimpleNamespace(zero_grad=lambda: None), + opt_param_scheduler=None, + config=config, + forward_backward_func=forward_backward, + iteration=0, + ) + + assert events == ["forward", ("setter", 450, 62500)] + with mock.patch.object(torch.distributed, "all_reduce") as all_reduce: + assert training_mod.consume_seqlen_stats_in_iteration() == (450, 62500) + all_reduce.assert_not_called() + finally: + training_mod._seqlen_stats_in_iteration = None + training_mod._seqlen_stats_active = False + training_mod._seqlen_stats_are_global = False + + +def test_training_log_uses_scheduled_microbatch_count_for_mtp(): + args = SimpleNamespace( + timing_log_level=0, + perform_rl_step=False, + micro_batch_size=1, + data_parallel_size=1, + world_size=1, + seq_length=8, + freeze_all_layers=False, + num_experts=None, + mtp_num_layers=1, + dsa_indexer_loss_coeff=None, + log_interval=100, + ) + + with ( + mock.patch.object(training_mod, "get_args", return_value=args), + mock.patch.object(training_mod, "get_timers", return_value=mock.MagicMock()), + mock.patch.object(training_mod, "get_tensorboard_writer", return_value=None), + mock.patch.object(training_mod, "get_wandb_writer", return_value=None), + mock.patch.object(training_mod, "get_one_logger", return_value=None), + mock.patch.object(training_mod, "get_energy_monitor", return_value=None), + mock.patch.object(training_mod, "get_num_microbatches", return_value=8), + mock.patch.object( + training_mod, "reduce_max_stat_across_model_parallel_group", return_value=None + ), + mock.patch.object(training_mod.one_logger_utils, "track_app_tag"), + mock.patch.object( + training_mod.MTPLossLoggingHelper, "track_mtp_metrics" + ) as track_mtp_metrics, + ): + training_mod.training_log( + loss_dict={}, + total_loss_dict={}, + learning_rate=None, + iteration=1, + loss_scale=1.0, + report_memory_flag=False, + skipped_iter=0, + grad_norm=None, + params_norm=None, + num_zeros_in_grad=None, + max_attention_logit=None, + num_microbatches=3, + ) + + assert track_mtp_metrics.call_args.args[0] == 1 / 3 + + +def test_training_log_uses_full_hybrid_moe_layout_including_mtp(): + args = SimpleNamespace( + timing_log_level=0, + perform_rl_step=False, + micro_batch_size=1, + data_parallel_size=1, + world_size=1, + seq_length=8, + freeze_all_layers=False, + num_experts=8, + moe_router_load_balancing_type=["seq_aux_loss"], + moe_z_loss_coeff=None, + moe_per_layer_logging=True, + moe_layer_freq=1, + hybrid_layer_pattern="CECE/WE", + mtp_num_layers=1, + dsa_indexer_loss_coeff=None, + log_interval=100, + ) + tracker = mock.MagicMock() + tracker.report.return_value = "" + + with ( + mock.patch.object(training_mod, "get_args", return_value=args), + mock.patch.object(training_mod, "get_timers", return_value=mock.MagicMock()), + mock.patch.object(training_mod, "get_tensorboard_writer", return_value=None), + mock.patch.object(training_mod, "get_wandb_writer", return_value=None), + mock.patch.object(training_mod, "get_one_logger", return_value=None), + mock.patch.object(training_mod, "get_energy_monitor", return_value=None), + mock.patch.object(training_mod, "get_num_microbatches", return_value=8), + mock.patch.object( + training_mod, "reduce_max_stat_across_model_parallel_group", return_value=None + ), + mock.patch.object(training_mod.one_logger_utils, "track_app_tag"), + mock.patch.object(training_mod, "get_moe_metrics_tracker", return_value=tracker), + mock.patch.object(training_mod.MTPLossLoggingHelper, "track_mtp_metrics"), + ): + training_mod.training_log( + loss_dict={}, + total_loss_dict={}, + learning_rate=None, + iteration=1, + loss_scale=1.0, + report_memory_flag=False, + skipped_iter=0, + grad_norm=None, + params_norm=None, + num_zeros_in_grad=None, + max_attention_logit=None, + ) + + assert tracker.report.call_args.kwargs["num_layers"] == 5 + assert tracker.report.call_args.kwargs["moe_layer_freq"] == [0, 1, 0, 1, 1] + assert tracker.report.call_args.kwargs["mtp_num_layers"] is None + + +def test_training_log_preserves_indexer_groups_only_for_te_graphs(): + for cuda_graph_impl, expected_preserve_groups in ( + ("none", False), + ("local", False), + ("full_iteration", False), + ("transformer_engine", True), + ): + args = SimpleNamespace( + timing_log_level=0, + perform_rl_step=False, + micro_batch_size=1, + data_parallel_size=1, + world_size=1, + seq_length=8, + freeze_all_layers=False, + num_experts=None, + mtp_num_layers=None, + dsa_indexer_loss_coeff=1.0, + num_layers=4, + csa_compress_ratios=[0, 4, 128, 0], + cuda_graph_impl=cuda_graph_impl, + log_interval=100, + ) + + with ( + mock.patch.object(training_mod, "get_args", return_value=args), + mock.patch.object(training_mod, "get_timers", return_value=mock.MagicMock()), + mock.patch.object(training_mod, "get_tensorboard_writer", return_value=None), + mock.patch.object(training_mod, "get_wandb_writer", return_value=None), + mock.patch.object(training_mod, "get_one_logger", return_value=None), + mock.patch.object(training_mod, "get_energy_monitor", return_value=None), + mock.patch.object(training_mod, "get_num_microbatches", return_value=8), + mock.patch.object( + training_mod, "reduce_max_stat_across_model_parallel_group", return_value=None + ), + mock.patch.object(training_mod.one_logger_utils, "track_app_tag"), + mock.patch.object( + training_mod.DSAIndexerLossLoggingHelper, "track_indexer_metrics" + ) as track_indexer_metrics, + ): + training_mod.training_log( + loss_dict={}, + total_loss_dict={}, + learning_rate=None, + iteration=1, + loss_scale=1.0, + report_memory_flag=False, + skipped_iter=0, + grad_norm=None, + params_norm=None, + num_zeros_in_grad=None, + max_attention_logit=None, + pg_collection=ProcessGroupCollection(), + ) + + assert track_indexer_metrics.call_args.kwargs["loss_scale"] == 1 / 8 + assert ( + track_indexer_metrics.call_args.kwargs["preserve_groups"] + is expected_preserve_groups + ) + + +def test_training_log_uses_scheduled_microbatch_count_for_dsa(): + args = SimpleNamespace( + timing_log_level=0, + perform_rl_step=False, + micro_batch_size=1, + data_parallel_size=1, + world_size=1, + seq_length=8, + freeze_all_layers=False, + num_experts=None, + mtp_num_layers=None, + dsa_indexer_loss_coeff=1.0, + num_layers=4, + csa_compress_ratios=[0, 4, 128, 0], + cuda_graph_impl="none", + log_interval=100, + ) + + with ( + mock.patch.object(training_mod, "get_args", return_value=args), + mock.patch.object(training_mod, "get_timers", return_value=mock.MagicMock()), + mock.patch.object(training_mod, "get_tensorboard_writer", return_value=None), + mock.patch.object(training_mod, "get_wandb_writer", return_value=None), + mock.patch.object(training_mod, "get_one_logger", return_value=None), + mock.patch.object(training_mod, "get_energy_monitor", return_value=None), + mock.patch.object(training_mod, "get_num_microbatches", return_value=8), + mock.patch.object( + training_mod, "reduce_max_stat_across_model_parallel_group", return_value=None + ), + mock.patch.object(training_mod.one_logger_utils, "track_app_tag"), + mock.patch.object( + training_mod.DSAIndexerLossLoggingHelper, "track_indexer_metrics" + ) as track_indexer_metrics, + ): + training_mod.training_log( + loss_dict={}, + total_loss_dict={}, + learning_rate=None, + iteration=1, + loss_scale=1.0, + report_memory_flag=False, + skipped_iter=0, + grad_norm=None, + params_norm=None, + num_zeros_in_grad=None, + max_attention_logit=None, + pg_collection=ProcessGroupCollection(), + num_microbatches=3, + ) + + assert track_indexer_metrics.call_args.kwargs["loss_scale"] == 1 / 3 diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py new file mode 100644 index 00000000000..ed4e4863c0b --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py @@ -0,0 +1,2288 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from unittest.mock import patch + +import pytest +import torch + +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.experimental_attention_variant.csa import ( + CompressedSparseAttention, + CompressedSparseAttentionSubmodules, + Compressor, + CompressorSubmodules, + CSAIndexer, + CSAIndexerSubmodules, + _apply_rope, + build_cu_seqlens_kv_full, + cat_per_segment, + get_compress_topk_idxs, + get_compress_topk_idxs_thd, + get_window_topk_idxs, + get_window_topk_idxs_thd, + unfused_compressed_sparse_attn, +) +from megatron.core.transformer.transformer_config import MLATransformerConfig +from tests.unit_tests.test_utilities import Utils + +try: + from fast_hadamard_transform import hadamard_transform as _hadamard_transform + + HAVE_HADAMARD = True +except ImportError: + HAVE_HADAMARD = False + _hadamard_transform = None + + +def mock_hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: + """Mock implementation of hadamard_transform for testing without the library installed.""" + return x * scale + + +class _DisabledContextTracker: + """Track whether a projection runs inside the FP8-disabled context.""" + + def __init__(self): + self.depth = 0 + self.entries = 0 + + def __call__(self, _config, is_init=False): + assert not is_init + return self + + def __enter__(self): + self.depth += 1 + self.entries += 1 + return self + + def __exit__(self, _exc_type, _exc_value, _traceback): + self.depth -= 1 + return False + + +@pytest.fixture(autouse=True) +def patch_hadamard_if_needed(): + """Automatically patch hadamard_transform in both dsa and csa modules if not installed.""" + if not HAVE_HADAMARD: + with ( + patch( + 'megatron.core.transformer.experimental_attention_variant.dsa.hadamard_transform', + mock_hadamard_transform, + ), + patch( + 'megatron.core.transformer.experimental_attention_variant.csa.rotate_activation', + lambda x: x * (x.size(-1) ** -0.5), + ), + ): + yield + else: + yield + + +# =========================================================================== +# Helper function tests +# =========================================================================== + + +class TestGetWindowTopkIdxs: + """Test get_window_topk_idxs helper.""" + + def test_basic_shape(self): + batch_size, seqlen, window_size = 2, 16, 4 + idxs = get_window_topk_idxs(window_size, batch_size, seqlen, torch.device("cpu")) + assert idxs.shape == (batch_size, seqlen, window_size) + + def test_causal_no_future(self): + """Indices should never exceed the query position.""" + seqlen, window_size = 32, 8 + idxs = get_window_topk_idxs(window_size, 1, seqlen, torch.device("cpu")) + for i in range(seqlen): + valid = idxs[0, i][idxs[0, i] >= 0] + assert torch.all(valid <= i), f"Position {i} has future indices" + + def test_invalid_marked_minus_one(self): + """Early positions that cannot fill the window should use -1.""" + seqlen, window_size = 8, 4 + idxs = get_window_topk_idxs(window_size, 1, seqlen, torch.device("cpu")) + assert idxs[0, 0, 0] == -1 or idxs[0, 0, 0] == 0 + for pos in range(window_size, seqlen): + assert torch.all(idxs[0, pos] >= 0), f"Position {pos} has invalid -1" + + def test_window_larger_than_seqlen(self): + """Window larger than sequence length should still work.""" + seqlen, window_size = 4, 16 + idxs = get_window_topk_idxs(window_size, 1, seqlen, torch.device("cpu")) + assert idxs.shape == (1, seqlen, window_size) + + +class TestGetCompressTopkIdxs: + """Test get_compress_topk_idxs helper.""" + + def test_basic_shape(self): + ratio, batch_size, seqlen, offset = 4, 2, 32, 32 + idxs = get_compress_topk_idxs(ratio, batch_size, seqlen, offset, torch.device("cpu")) + n_compressed = seqlen // ratio + assert idxs.shape == (batch_size, seqlen, n_compressed) + + def test_offset_applied(self): + """Valid indices should be >= offset.""" + ratio, seqlen, offset = 4, 32, 100 + idxs = get_compress_topk_idxs(ratio, 1, seqlen, offset, torch.device("cpu")) + valid = idxs[idxs >= 0] + if valid.numel() > 0: + assert torch.all(valid >= offset), "Valid indices should be offset" + + def test_causal_no_future(self): + """Compressed indices should respect causality.""" + ratio, seqlen, offset = 4, 32, 32 + idxs = get_compress_topk_idxs(ratio, 1, seqlen, offset, torch.device("cpu")) + for i in range(seqlen): + n_valid = (i + 1) // ratio + valid = idxs[0, i][idxs[0, i] >= 0] + assert valid.numel() <= n_valid, f"Position {i} has too many valid compressed indices" + + def test_ratio_128(self): + """Test with large compression ratio.""" + ratio, seqlen, offset = 128, 256, 256 + idxs = get_compress_topk_idxs(ratio, 1, seqlen, offset, torch.device("cpu")) + assert idxs.shape == (1, seqlen, seqlen // ratio) + + +# =========================================================================== +# unfused_compressed_sparse_attn tests +# =========================================================================== + + +class TestUnfusedCompressedSparseAttn: + """Test the unfused compressed sparse attention kernel.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + yield + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_output_shape(self): + """Test output shape of unfused compressed sparse attention.""" + sq, b, np_, hn = 16, 2, 4, 64 + n_kv = sq + sq // 4 + topk = 8 + + query = torch.randn(sq, b, np_, hn, dtype=torch.bfloat16).cuda() + kv_full = torch.randn(n_kv, b, hn, dtype=torch.bfloat16).cuda() + attn_sink = torch.zeros(np_, dtype=torch.float32).cuda() + topk_indices = torch.randint(0, n_kv, (b, sq, topk), dtype=torch.int32).cuda() + softmax_scale = hn**-0.5 + + output = unfused_compressed_sparse_attn( + query, kv_full, attn_sink, topk_indices, softmax_scale + ) + + assert output.shape == (sq, b, np_ * hn) + assert output.dtype == query.dtype + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_invalid_indices_masked(self): + """Test that -1 indices are properly masked.""" + sq, b, np_, hn = 8, 1, 2, 32 + n_kv = sq + topk = 4 + + query = torch.randn(sq, b, np_, hn, dtype=torch.bfloat16).cuda() + kv_full = torch.randn(n_kv, b, hn, dtype=torch.bfloat16).cuda() + attn_sink = torch.zeros(np_, dtype=torch.float32).cuda() + + topk_indices = torch.full((b, sq, topk), -1, dtype=torch.int32).cuda() + topk_indices[:, :, 0] = 0 + softmax_scale = hn**-0.5 + + output = unfused_compressed_sparse_attn( + query, kv_full, attn_sink, topk_indices, softmax_scale + ) + assert not torch.isnan(output).any(), "Output should not contain NaN" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_gradient_flow(self): + """Test that gradients flow through sparse attention.""" + sq, b, np_, hn = 8, 1, 2, 32 + n_kv = sq + topk = 4 + + query = torch.randn(sq, b, np_, hn, dtype=torch.float32).cuda().requires_grad_(True) + kv_full = torch.randn(n_kv, b, hn, dtype=torch.float32).cuda().requires_grad_(True) + attn_sink = torch.nn.Parameter(torch.zeros(np_, dtype=torch.float32).cuda()) + + topk_indices = torch.randint(0, n_kv, (b, sq, topk), dtype=torch.int32).cuda() + softmax_scale = hn**-0.5 + + output = unfused_compressed_sparse_attn( + query, kv_full, attn_sink, topk_indices, softmax_scale + ) + loss = output.sum() + loss.backward() + + assert query.grad is not None + assert kv_full.grad is not None + assert attn_sink.grad is not None + + +# =========================================================================== +# Compressor tests +# =========================================================================== + + +def _make_mla_config( + num_layers=4, + hidden_size=256, + num_attention_heads=16, + v_head_dim=64, + qk_pos_emb_head_dim=32, + csa_compress_ratios=None, + csa_window_size=8, + csa_dense_mode=False, + tensor_model_parallel_size=1, + sequence_parallel=False, + dsa_indexer_n_heads=8, + dsa_indexer_head_dim=64, + dsa_indexer_topk=8, + dsa_indexer_loss_coeff=0.0, + dsa_indexer_use_sparse_loss=False, + rope_type='rope', +): + """Helper to create MLATransformerConfig for CSA tests.""" + if csa_compress_ratios is None: + csa_compress_ratios = [0] * num_layers + return MLATransformerConfig( + num_layers=num_layers, + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + use_cpu_initialization=True, + bf16=True, + params_dtype=torch.bfloat16, + tensor_model_parallel_size=tensor_model_parallel_size, + sequence_parallel=sequence_parallel, + q_lora_rank=64, + kv_lora_rank=64, + qk_head_dim=v_head_dim - qk_pos_emb_head_dim, + qk_pos_emb_head_dim=qk_pos_emb_head_dim, + v_head_dim=v_head_dim, + rope_type=rope_type, + rotary_base=10000, + rotary_percent=1.0, + multi_latent_attention=True, + csa_compress_ratios=csa_compress_ratios, + csa_window_size=csa_window_size, + csa_dense_mode=csa_dense_mode, + dsa_indexer_n_heads=dsa_indexer_n_heads, + dsa_indexer_head_dim=dsa_indexer_head_dim, + dsa_indexer_topk=dsa_indexer_topk, + dsa_indexer_loss_coeff=dsa_indexer_loss_coeff, + dsa_indexer_use_sparse_loss=dsa_indexer_use_sparse_loss, + ) + + +def _make_compressor_submodules(): + """Create Compressor submodules spec.""" + from megatron.core.extensions.transformer_engine import TELinear, TENorm + from megatron.core.transformer.spec_utils import ModuleSpec + + return CompressorSubmodules( + linear_wkv=ModuleSpec(module=TELinear), + linear_wgate=ModuleSpec(module=TELinear), + norm=ModuleSpec(module=TENorm), + ) + + +def _make_csa_indexer_submodules(): + """Create CSAIndexer submodules spec.""" + from megatron.core.extensions.transformer_engine import TELinear, TENorm + from megatron.core.transformer.spec_utils import ModuleSpec + + return CSAIndexerSubmodules( + linear_wq_b=ModuleSpec(module=TELinear), + linear_weights_proj=ModuleSpec(module=TELinear), + compressor=ModuleSpec(module=Compressor, submodules=_make_compressor_submodules()), + ) + + +def _make_csa_submodules(): + """Create CompressedSparseAttention submodules spec.""" + from megatron.core.transformer.spec_utils import ModuleSpec + + return CompressedSparseAttentionSubmodules( + compressor=ModuleSpec(module=Compressor, submodules=_make_compressor_submodules()), + indexer=ModuleSpec(module=CSAIndexer, submodules=_make_csa_indexer_submodules()), + ) + + +# =========================================================================== +# Compressor tests +# =========================================================================== + + +@pytest.mark.parametrize("compress_ratio", [4, 128]) +class TestCompressor: + """Test Compressor module.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config(csa_compress_ratios=[4, 128, 4, 128]) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_compressor_output_shape(self, compress_ratio): + """Test that compressor produces correct output shape.""" + seq_len = 256 + batch_size = 2 + head_dim = self.config.v_head_dim + + compressor = Compressor( + config=self.config, + submodules=_make_compressor_submodules(), + compress_ratio=compress_ratio, + head_dim=head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + output = compressor(x) + + expected_len = seq_len // compress_ratio + assert output is not None + assert output.shape == (expected_len, batch_size, head_dim) + assert output.dtype == torch.bfloat16 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_compressor_too_short_input(self, compress_ratio): + """Test that compressor returns None when input is shorter than compress_ratio.""" + short_len = compress_ratio - 1 + batch_size = 2 + head_dim = self.config.v_head_dim + + compressor = Compressor( + config=self.config, + submodules=_make_compressor_submodules(), + compress_ratio=compress_ratio, + head_dim=head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + + x = torch.randn(short_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + output = compressor(x) + assert output is None + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_compressor_gradient_flow(self, compress_ratio): + """Test that gradients flow through the compressor.""" + seq_len = 256 + batch_size = 2 + head_dim = self.config.v_head_dim + + compressor = Compressor( + config=self.config, + submodules=_make_compressor_submodules(), + compress_ratio=compress_ratio, + head_dim=head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + + x = ( + torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) + output = compressor(x) + loss = output.sum() + loss.backward() + + assert x.grad is not None + for name, param in compressor.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"Parameter {name} has no gradient" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_projection_disables_fp8(self, compress_ratio, monkeypatch): + compressor = Compressor( + config=self.config, + submodules=_make_compressor_submodules(), + compress_ratio=compress_ratio, + head_dim=self.config.v_head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + tracker = _DisabledContextTracker() + calls = [] + + for name, projection in ( + ('linear_wkv', compressor.linear_wkv), + ('linear_wgate', compressor.linear_wgate), + ): + original_forward = projection.forward + + def checked_forward(*args, _name=name, _forward=original_forward, **kwargs): + assert tracker.depth > 0, f"{_name} ran outside the FP8-disabled context" + calls.append(_name) + return _forward(*args, **kwargs) + + monkeypatch.setattr(projection, 'forward', checked_forward) + + monkeypatch.setattr( + 'megatron.core.transformer.experimental_attention_variant.csa.get_fp8_disabled_context', + tracker, + ) + x = torch.randn( + compress_ratio * 2, + 1, + self.config.hidden_size, + dtype=torch.bfloat16, + device='cuda', + ) + compressor(x) + + assert calls == ['linear_wkv', 'linear_wgate'] + assert tracker.entries == 1 + + +# =========================================================================== +# CSAIndexer tests +# =========================================================================== + + +@pytest.mark.parametrize("seqlen", [32, 128]) +class TestCSAIndexer: + """Test CSAIndexer module basic functionality.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.compress_ratio = 4 + cls.config = _make_mla_config(csa_compress_ratios=[4, 4, 4, 4], dsa_indexer_topk=8) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + cls.indexer = CSAIndexer( + config=cls.config, + submodules=_make_csa_indexer_submodules(), + compress_ratio=cls.compress_ratio, + rotary_pos_emb=cls.rotary_pos_emb, + pg_collection=cls.pg_collection, + ) + + yield + Utils.destroy_model_parallel() + + def test_csa_indexer_constructor(self, seqlen): + """Test CSAIndexer initialization.""" + assert isinstance(self.indexer, CSAIndexer) + assert self.indexer.compress_ratio == self.compress_ratio + assert self.indexer.index_n_heads == self.config.dsa_indexer_n_heads + assert self.indexer.index_head_dim == self.config.dsa_indexer_head_dim + assert self.indexer.index_topk == self.config.dsa_indexer_topk + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_csa_indexer_forward(self, seqlen): + """Test CSAIndexer forward pass.""" + batch_size = 2 + self.indexer.cuda() + + x = torch.randn(seqlen, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seqlen, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + index_scores, topk_indices = self.indexer(x, qr) + n_compressed = seqlen // self.compress_ratio + effective_topk = min(self.config.dsa_indexer_topk, n_compressed) + + assert index_scores.shape == (batch_size, seqlen, n_compressed) + assert topk_indices.shape == (batch_size, seqlen, effective_topk) + assert index_scores.dtype == torch.float32 + assert topk_indices.dtype == torch.long + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_csa_indexer_forward_before_topk(self, seqlen): + """Test CSAIndexer forward_before_topk.""" + batch_size = 2 + self.indexer.cuda() + + x = torch.randn(seqlen, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seqlen, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + q, k, weights = self.indexer.forward_before_topk(x, qr) + + assert q.shape == ( + seqlen, + batch_size, + self.config.dsa_indexer_n_heads, + self.config.dsa_indexer_head_dim, + ) + n_compressed = seqlen // self.compress_ratio + assert k.shape == (n_compressed, batch_size, self.config.dsa_indexer_head_dim) + assert weights.shape == (seqlen, batch_size, self.config.dsa_indexer_n_heads) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_weights_projection_disables_fp8(self, seqlen, monkeypatch): + tracker = _DisabledContextTracker() + self.indexer.cuda() + original_forward = self.indexer.linear_weights_proj.forward + + def checked_forward(*args, **kwargs): + assert tracker.depth > 0, "indexer weights projection ran under FP8" + return original_forward(*args, **kwargs) + + monkeypatch.setattr(self.indexer.linear_weights_proj, 'forward', checked_forward) + monkeypatch.setattr( + 'megatron.core.transformer.experimental_attention_variant.csa.get_fp8_disabled_context', + tracker, + ) + x = torch.randn( + seqlen, + 1, + self.config.hidden_size, + dtype=torch.bfloat16, + device='cuda', + ) + weights = self.indexer._project_weights(x) + + assert weights.shape == (seqlen, 1, self.config.dsa_indexer_n_heads) + assert tracker.entries == 1 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_csa_indexer_with_mask(self, seqlen): + """Test CSAIndexer with causal mask.""" + batch_size = 2 + self.indexer.cuda() + + x = torch.randn(seqlen, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seqlen, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + n_compressed = seqlen // self.compress_ratio + causal_mask = torch.arange(n_compressed, device=x.device).unsqueeze(0).expand(seqlen, -1) + positions = torch.arange(1, seqlen + 1, device=x.device).unsqueeze(1) + causal_mask = ( + torch.where(causal_mask >= positions // self.compress_ratio, float("-inf"), 0.0) + .unsqueeze(0) + .expand(batch_size, -1, -1) + ) + + index_scores, topk_indices = self.indexer(x, qr, mask=causal_mask) + + effective_topk = min(self.config.dsa_indexer_topk, n_compressed) + assert index_scores.shape == (batch_size, seqlen, n_compressed) + assert topk_indices.shape == (batch_size, seqlen, effective_topk) + + +# =========================================================================== +# CompressedSparseAttention tests +# =========================================================================== + + +class TestCompressedSparseAttentionRatio1: + """Test CompressedSparseAttention with compress_ratio=1 (window-only).""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config(csa_compress_ratios=[0, 0, 0, 0], csa_window_size=8) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + cls.csa = CompressedSparseAttention( + config=cls.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=cls.pg_collection, + rotary_pos_emb=rotary_pos_emb, + compress_ratio=0, + ) + + yield + Utils.destroy_model_parallel() + + def test_ratio1_no_compressor(self): + """With ratio=1, compressor and indexer should not be built.""" + assert self.csa.compressor is None + assert self.csa.indexer is None + + def test_mtp_layer_number_is_offset(self): + """MTP attention layers are numbered after all decoder layers.""" + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + compress_ratio=0, + is_mtp_layer=True, + ) + + assert csa.layer_number == self.config.num_layers + 1 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_ratio1_forward(self): + """Test forward pass with window-only attention.""" + seq_len = 32 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + self.csa.cuda() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.bfloat16).cuda() + value = key.clone() + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = self.csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + + assert output.shape == (seq_len, batch_size, np_ * hn) + assert output.dtype == torch.bfloat16 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_ratio1_backward(self): + """Test backward pass with window-only attention.""" + seq_len = 32 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + self.csa.train() + self.csa.cuda() + + query = ( + torch.randn(seq_len, batch_size, np_, hn, dtype=torch.float32) + .cuda() + .requires_grad_(True) + ) + key = ( + torch.randn(seq_len, batch_size, 1, hn, dtype=torch.float32).cuda().requires_grad_(True) + ) + value = key.clone().detach().requires_grad_(True) + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = self.csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + loss = output.sum() + loss.backward() + + assert query.grad is not None + assert key.grad is not None + + +@pytest.mark.parametrize("compress_ratio", [4, 128]) +class TestCompressedSparseAttentionCompressed: + """Test CompressedSparseAttention with compress_ratio > 1.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config( + csa_compress_ratios=[4, 128, 4, 128], + csa_window_size=8, + dsa_indexer_topk=8, + dsa_indexer_loss_coeff=1.0, + ) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + def _get_layer_number(self, compress_ratio): + """Return a layer_number (1-indexed) whose compress_ratio matches.""" + for i, r in enumerate(self.config.csa_compress_ratios): + if r == compress_ratio: + return i + 1 + raise ValueError(f"No layer with compress_ratio={compress_ratio}") + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_constructor(self, compress_ratio): + """Test that compressor/indexer are conditionally built.""" + layer_number = self._get_layer_number(compress_ratio) + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, + ).cuda() + + assert csa.compressor is not None + if compress_ratio == 4: + assert csa.indexer is not None + elif compress_ratio == 128: + assert csa.indexer is None + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_forward(self, compress_ratio): + """Test forward pass with compressed attention.""" + seq_len = 256 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + layer_number = self._get_layer_number(compress_ratio) + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, + ).cuda() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.bfloat16).cuda() + value = key.clone() + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + + assert output.shape == (seq_len, batch_size, np_ * hn) + assert not torch.isnan(output).any() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_backward(self, compress_ratio): + """Test backward pass with compressed attention.""" + seq_len = 256 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + layer_number = self._get_layer_number(compress_ratio) + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, + ).cuda() + csa.train() + + query = ( + torch.randn(seq_len, batch_size, np_, hn, dtype=torch.float32) + .cuda() + .requires_grad_(True) + ) + key = ( + torch.randn(seq_len, batch_size, 1, hn, dtype=torch.float32).cuda().requires_grad_(True) + ) + value = key.clone().detach().requires_grad_(True) + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + loss = output.sum() + loss.backward() + + assert query.grad is not None + assert key.grad is not None + + for name, param in csa.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"Parameter {name} has no gradient" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_eval_mode(self, compress_ratio): + """Test forward pass in eval mode.""" + seq_len = 256 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + layer_number = self._get_layer_number(compress_ratio) + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, + ).cuda() + csa.eval() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.bfloat16).cuda() + value = key.clone() + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + with torch.no_grad(): + output = csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + + assert output.shape == (seq_len, batch_size, np_ * hn) + assert not torch.isnan(output).any() + + +# =========================================================================== +# _apply_rope tests +# =========================================================================== + + +class TestApplyRope: + """Test ``_apply_rope`` — the layout-aware RoPE wrapper used by + Compressor / CSAIndexer / hybrid-attention callers. + + Behaviours covered: + + * 3-D ``[seq, batch, head_dim]`` and 4-D ``[seq, batch, heads, head_dim]`` + inputs both work (3-D gets a temporary head-dim unsqueeze). + * Only the trailing ``pos_dim`` components are rotated; the leading + ``nope_dim`` slice is bit-exact unchanged. + * Both ``RotaryEmbedding`` (returns ``Tensor``) and + ``YarnRotaryEmbedding`` (returns ``(emb, mscale)`` tuple) — DSv4 + hybrid silently swaps the class based on ``compress_ratio``. + * Both unfused and fused (``config.apply_rope_fusion=True``) paths + produce the same output (within bf16 precision). + * For ``ratio > 1`` the rotary table is built at + ``rotary_seq_len * ratio`` and strided by ``ratio``. + """ + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(0) + model_parallel_cuda_manual_seed(0) + cls = request.cls + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + # head_dim 32 = nope 24 + pos 8 + cls.config = _make_mla_config(v_head_dim=32, qk_pos_emb_head_dim=8) + yield + Utils.destroy_model_parallel() + + def _make_rotary(self, kind: str): + from megatron.core.models.common.embeddings import RotaryEmbedding, YarnRotaryEmbedding + + pos_dim = self.config.qk_pos_emb_head_dim + if kind == 'rope': + return RotaryEmbedding( + pos_dim, rotary_percent=1.0, rotary_base=10000, cp_group=self.pg_collection.cp + ) + if kind == 'yarn': + return YarnRotaryEmbedding( + pos_dim, + rotary_base=40000, + scaling_factor=40, + original_max_position_embeddings=4096, + beta_fast=32, + beta_slow=1, + mscale=1.0, + mscale_all_dim=0.0, + cp_group=self.pg_collection.cp, + ) + raise ValueError(kind) + + def _config_with(self, *, apply_rope_fusion: bool): + # Reuse the class-level config; only flip the fusion flag. + cfg = self.config + cfg.apply_rope_fusion = apply_rope_fusion + return cfg + + _ROTARY_FUSION_COMBOS = [ + pytest.param('rope', False, id='rope-unfused'), + pytest.param('rope', True, id='rope-fused'), + pytest.param('yarn', False, id='yarn-unfused'), + pytest.param('yarn', True, id='yarn-fused'), + ] + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize(("rotary_kind", "apply_rope_fusion"), _ROTARY_FUSION_COMBOS) + @pytest.mark.parametrize("input_ndim", [3, 4], ids=['3d', '4d']) + @pytest.mark.parametrize("ratio", [1, 4], ids=['ratio_1', 'ratio_4']) + def test_apply_rope(self, rotary_kind, apply_rope_fusion, input_ndim, ratio): + """Output shape == input shape; no NaN; nope-dim slice is + bit-exact unchanged. Sweeps the valid combinations of rotary + class × apply_rope_fusion × input rank × ratio. Yarn's + tuple-return is covered by the ``'yarn-*'`` combos. + """ + rotary = self._make_rotary(rotary_kind).cuda() + nope = self.config.v_head_dim - self.config.qk_pos_emb_head_dim + pos = self.config.qk_pos_emb_head_dim + head_dim = nope + pos + seq, batch, heads = 8, 2, 4 + cfg = self._config_with(apply_rope_fusion=apply_rope_fusion) + + shape = (seq, batch, head_dim) if input_ndim == 3 else (seq, batch, heads, head_dim) + x = torch.randn(*shape, dtype=torch.bfloat16, device='cuda') + # ``fused_mla_rope_inplace`` mutates the input — give it a copy so + # the nope-dim equality check below still has the original. + out = _apply_rope( + x.clone(), + nope, + pos, + rotary, + cfg, + rotary_seq_len=seq, + ratio=ratio, + cp_group=self.pg_collection.cp, + ) + + assert out.shape == x.shape + assert out.dtype == x.dtype + assert not torch.isnan(out).any() + # The leading nope_dim slice is the identity portion of RoPE. + assert torch.equal( + out[..., :nope], x[..., :nope] + ), "RoPE must not touch the first nope_dim components" + # Trailing pos_dim should rotate at non-zero positions. + pe_changed = (out[..., nope:] != x[..., nope:]).any(dim=-1).flatten() + assert pe_changed[ + 1: + ].any(), "RoPE should rotate the trailing pos_dim components for seq > 0" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("rotary_kind", ['rope', 'yarn']) + def test_3d_input_matches_4d_with_single_head(self, rotary_kind): + """For a single-head input, the 3-D ``(s, b, d)`` and 4-D + ``(s, b, 1, d)`` invocations must produce numerically identical + output (3-D path just inserts a temporary head dim). + """ + rotary = self._make_rotary(rotary_kind).cuda() + nope = self.config.v_head_dim - self.config.qk_pos_emb_head_dim + pos = self.config.qk_pos_emb_head_dim + head_dim = nope + pos + seq, batch = 8, 2 + cfg = self._config_with(apply_rope_fusion=False) + + x_3d = torch.randn(seq, batch, head_dim, dtype=torch.bfloat16, device='cuda') + x_4d = x_3d.unsqueeze(-2) + + out_3d = _apply_rope( + x_3d, + nope, + pos, + rotary, + cfg, + rotary_seq_len=seq, + ratio=1, + cp_group=self.pg_collection.cp, + ) + out_4d = _apply_rope( + x_4d, + nope, + pos, + rotary, + cfg, + rotary_seq_len=seq, + ratio=1, + cp_group=self.pg_collection.cp, + ) + + assert out_3d.shape == x_3d.shape + assert out_4d.shape == x_4d.shape + assert torch.equal(out_3d, out_4d.squeeze(-2)) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("rotary_kind", ['rope', 'yarn']) + def test_ratio_strides_rotary_table(self, rotary_kind): + """For ``ratio > 1``, the rotary table is built at + ``rotary_seq_len * ratio`` and strided by ``ratio``. The result + with ``ratio=k`` must equal an ``apply_rope`` call on the same + positions of a length-``rotary_seq_len * k`` table. + """ + rotary = self._make_rotary(rotary_kind).cuda() + nope = self.config.v_head_dim - self.config.qk_pos_emb_head_dim + pos = self.config.qk_pos_emb_head_dim + head_dim = nope + pos + seq, batch, heads, ratio = 4, 1, 2, 4 + cfg = self._config_with(apply_rope_fusion=False) + + x_comp = torch.randn(seq, batch, heads, head_dim, dtype=torch.bfloat16, device='cuda') + out_comp = _apply_rope( + x_comp.clone(), + nope, + pos, + rotary, + cfg, + rotary_seq_len=seq, + ratio=ratio, + cp_group=self.pg_collection.cp, + ) + + x_full = torch.zeros( + seq * ratio, batch, heads, head_dim, dtype=torch.bfloat16, device='cuda' + ) + x_full[::ratio][:seq] = x_comp + out_full = _apply_rope( + x_full, + nope, + pos, + rotary, + cfg, + rotary_seq_len=seq * ratio, + ratio=1, + cp_group=self.pg_collection.cp, + ) + out_ref = out_full[::ratio][:seq] + + assert torch.allclose(out_comp, out_ref, rtol=1e-3, atol=1e-3), ( + f"ratio={ratio} stride mismatch: " + f"max abs diff = {(out_comp - out_ref).abs().max().item():.3e}" + ) + + +# =========================================================================== +# csa_dense_mode tests +# =========================================================================== + + +class TestCompressedSparseAttentionDenseMode: + """Test that csa_dense_mode=True disables the indexer for ratio=4 layers.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config( + csa_compress_ratios=[4, 128, 4, 128], csa_window_size=8, csa_dense_mode=True + ) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dense_mode_disables_indexer_for_ratio4(self): + """With csa_dense_mode=True, ratio=4 layers should NOT build an indexer.""" + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=4, + ).cuda() + + assert csa.compress_ratio == 4 + assert csa.compressor is not None, "Compressor should still be built" + assert csa.indexer is None, "Indexer should be disabled in dense mode" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dense_mode_forward_ratio4(self): + """Forward pass should work for ratio=4 in dense mode (uses all compressed positions).""" + seq_len = 256 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=4, + ).cuda() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.bfloat16).cuda() + value = key.clone() + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + + assert output.shape == (seq_len, batch_size, np_ * hn) + assert not torch.isnan(output).any() + +# =========================================================================== +# THD packed-sequence helpers +# =========================================================================== + + +def _cu_seqlens(seg_lens, device='cpu'): + """``(B+1,)`` int32 cu_seqlens from a list of per-segment lengths.""" + return torch.tensor( + [0] + list(torch.tensor(seg_lens, dtype=torch.int64).cumsum(0).tolist()), + dtype=torch.int32, + device=device, + ) + + +class TestCsaThdIndexHelpers: + """CSA THD index helpers — pure-Python, no GPU. Mirrors the + organisation of ``TestThdPureHelpers`` in ``test_dsa_kernels.py``: + one mega-class with section comments per helper, since each helper + only needs 2–3 tests and they share no fixtures. + + Helpers covered: + + * ``get_window_topk_idxs_thd`` — per-segment sliding window. + * ``get_compress_topk_idxs_thd`` — per-segment all-compressed + indices shifted to full-KV space. + * ``build_cu_seqlens_kv_full`` — per-segment lens of the + ``[kv, compressed_kv]`` concat. + * ``cat_per_segment`` — per-segment concat into the + THD-packed full-KV layout. + """ + + # ---- get_window_topk_idxs_thd -------------------------------------- + + def test_window_shape_dtype_and_local_indices(self): + """Window indices are LOCAL within each segment — they reset to 0 + at each segment boundary (not global flat KV ids). + """ + cu = _cu_seqlens([4, 3]) # = [0, 4, 7] + out = get_window_topk_idxs_thd(window_size=3, cu_seqlens_q=cu) + assert out.shape == (7, 3) + assert out.dtype == torch.int32 + expected = torch.tensor( + [[0, -1, -1], [0, 1, -1], [0, 1, 2], [1, 2, 3], [0, -1, -1], [0, 1, -1], [0, 1, 2]], + dtype=torch.int32, + ) + assert torch.equal(out, expected) + + def test_window_causality_no_future(self): + """No window index should exceed the query's position-in-segment.""" + cu = _cu_seqlens([5, 6, 3]) + out = get_window_topk_idxs_thd(window_size=4, cu_seqlens_q=cu) + seq_lens = (cu[1:] - cu[:-1]).tolist() + offsets = cu[:-1].tolist() + for b, (offset, slen) in enumerate(zip(offsets, seq_lens)): + for s in range(slen): + row = out[offset + s] + valid = row[row >= 0] + assert (valid <= s).all(), f"seg {b}, pos {s}: window index exceeds position" + + # ---- get_compress_topk_idxs_thd ------------------------------------ + + @pytest.mark.parametrize( + "q_segs, kv_segs, comp_segs, expected_shape, expected_ranges", + [ + ([8, 4], [5, 3], [2, 1], (12, 2), {(0, 8): (5, 7), (8, 12): (3, 4)}), + ([3, 2], [3, 2], [0, 0], (5, 0), {}), + ], + ids=["multi_seg_offsets", "no_compressed_empty"], + ) + def test_compress_shape_and_offset( + self, q_segs, kv_segs, comp_segs, expected_shape, expected_ranges + ): + """Valid indices live in the correct per-segment range, or output is + empty when all segments are shorter than ratio. + """ + ratio = 4 + out = get_compress_topk_idxs_thd( + ratio, _cu_seqlens(q_segs), _cu_seqlens(kv_segs), _cu_seqlens(comp_segs) + ) + assert out.shape == expected_shape + for (start, end), (lo, hi) in expected_ranges.items(): + valid = out[start:end][out[start:end] >= 0] + assert (valid >= lo).all() and (valid < hi).all() + + def test_compress_causal_n_valid_per_pos(self): + """Per-row valid count == ``min(seqlen_compressed[b], (pos+1)//ratio)``.""" + ratio = 4 + out = get_compress_topk_idxs_thd( + ratio, _cu_seqlens([8]), _cu_seqlens([5]), _cu_seqlens([2]) + ) + for pos in range(8): + n_valid_expected = min(2, (pos + 1) // ratio) + n_valid_actual = int((out[pos] >= 0).sum()) + assert n_valid_actual == n_valid_expected, f"pos {pos}: count mismatch" + + # ---- build_cu_seqlens_kv_full -------------------------------------- + + def test_build_cu_seqlens_kv_full_basic(self): + cu_kv = _cu_seqlens([4, 3, 5]) + cu_comp = _cu_seqlens([1, 0, 2]) + out = build_cu_seqlens_kv_full(cu_kv, cu_comp) + # full lens = [4+1, 3+0, 5+2] = [5, 3, 7]; cumsum = [0, 5, 8, 15]. + assert out.tolist() == [0, 5, 8, 15] + assert out.dtype == cu_kv.dtype + + def test_build_cu_seqlens_kv_full_empty_compressed(self): + """When compressed is all zeros, full == kv.""" + cu_kv = _cu_seqlens([3, 4]) + cu_comp = _cu_seqlens([0, 0]) + out = build_cu_seqlens_kv_full(cu_kv, cu_comp) + assert torch.equal(out, cu_kv) + + # ---- cat_per_segment ------------------------------------------------ + + def test_cat_per_segment_basic_concat(self): + kv_lens = [3, 2] + comp_lens = [1, 2] + d = 2 + cu_kv = _cu_seqlens(kv_lens) + cu_comp = _cu_seqlens(comp_lens) + cu_full = build_cu_seqlens_kv_full(cu_kv, cu_comp) + + # Distinct values so we can verify each row's source. + kv = torch.arange(sum(kv_lens) * d, dtype=torch.float32).reshape(-1, d) + comp = (torch.arange(sum(comp_lens) * d, dtype=torch.float32) + 100).reshape(-1, d) + + out = cat_per_segment(kv, comp, cu_kv, cu_comp, cu_full) + assert out.shape == (sum(kv_lens) + sum(comp_lens), d) + # Segment 0: kv rows 0..2, then comp row 0. + assert torch.equal(out[0:3], kv[0:3]) + assert torch.equal(out[3:4], comp[0:1]) + # Segment 1: kv rows 3..4, then comp rows 1..2. + assert torch.equal(out[4:6], kv[3:5]) + assert torch.equal(out[6:8], comp[1:3]) + + def test_cat_per_segment_none_compressed_returns_kv(self): + """``compressed_kv_thd is None`` short-circuits to ``kv_thd``.""" + cu_kv = _cu_seqlens([3, 2]) + cu_comp = _cu_seqlens([0, 0]) + cu_full = build_cu_seqlens_kv_full(cu_kv, cu_comp) + kv = torch.randn(5, 2) + out = cat_per_segment(kv, None, cu_kv, cu_comp, cu_full) + assert out is kv + + +# =========================================================================== +# unfused_compressed_sparse_attn THD branch +# =========================================================================== + + +class TestUnfusedCompressedSparseAttnThd: + """``unfused_compressed_sparse_attn`` dispatches on ``query.ndim``: + 3-D selects the THD branch (flat layout, global topk ids). + """ + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + yield + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_output_shape(self): + """THD inputs (3-D query, 2-D kv) produce 2-D ``(total_q, np * hn)``.""" + total_q, np_, hn = 12, 4, 64 + total_kv = 24 + topk = 4 + + query = torch.randn(total_q, np_, hn, dtype=torch.bfloat16).cuda() + kv_full = torch.randn(total_kv, hn, dtype=torch.bfloat16).cuda() + attn_sink = torch.zeros(np_, dtype=torch.float32).cuda() + topk_indices = torch.randint(0, total_kv, (total_q, topk), dtype=torch.int32).cuda() + + out = unfused_compressed_sparse_attn(query, kv_full, attn_sink, topk_indices, hn**-0.5) + assert out.shape == (total_q, np_ * hn) + assert out.dtype == query.dtype + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_invalid_indices_masked(self): + """``-1`` indices in the THD topk should contribute 0 (no NaN).""" + total_q, np_, hn = 6, 2, 32 + total_kv = 8 + topk = 4 + + query = torch.randn(total_q, np_, hn, dtype=torch.bfloat16).cuda() + kv_full = torch.randn(total_kv, hn, dtype=torch.bfloat16).cuda() + attn_sink = torch.zeros(np_, dtype=torch.float32).cuda() + + topk_indices = torch.full((total_q, topk), -1, dtype=torch.int32).cuda() + topk_indices[:, 0] = 0 # one valid position per row + + out = unfused_compressed_sparse_attn(query, kv_full, attn_sink, topk_indices, hn**-0.5) + assert not torch.isnan(out).any() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_matches_sbhd_b1_equivalent(self): + """THD (3-D query) on a single-batch problem produces the same + per-token output as SBHD (4-D query) with ``b=1`` on the same + data — both should hit the shared core inlined into the function. + """ + sq, np_, hn = 8, 2, 32 + n_kv = 16 + topk = 4 + sm = hn**-0.5 + + torch.manual_seed(0) + # SBHD layout (b=1) and THD-equivalent (squeezed). + query_sbhd = torch.randn(sq, 1, np_, hn, dtype=torch.bfloat16).cuda() + kv_sbhd = torch.randn(n_kv, 1, hn, dtype=torch.bfloat16).cuda() + attn_sink = torch.zeros(np_, dtype=torch.float32).cuda() + + # SBHD topk: per-batch LOCAL ids in [0, n_kv). + topk_local = torch.randint(0, n_kv, (1, sq, topk), dtype=torch.int32).cuda() + + # THD topk: flat-global ids; for b=1 these match the local ids. + topk_global = topk_local.squeeze(0) + + out_sbhd = unfused_compressed_sparse_attn( + query_sbhd, kv_sbhd, attn_sink, topk_local, sm + ) # (sq, 1, np * hn) + out_thd = unfused_compressed_sparse_attn( + query_sbhd.squeeze(1), kv_sbhd.squeeze(1), attn_sink, topk_global, sm + ) # (sq, np * hn) + + # Same math, just different output layout. + assert torch.allclose(out_sbhd.squeeze(1), out_thd, atol=1e-3, rtol=1e-3) + + +# =========================================================================== +# THD: Compressor / CSAIndexer / CompressedSparseAttention integration +# =========================================================================== +# +# These integration tests exercise the THD branches of the full +# Compressor / CSAIndexer / CompressedSparseAttention modules — the +# layer above the kernel-level THD tests in test_dsa_kernels.py and the +# autograd-Function tests in test_attention_variant_dsa.py. +# +# Strategy: most tests use a B=1 single-segment THD input and compare +# against the same data run through the SBHD path with b=1. For B=1 +# the two layouts go through equivalent math (sparse-attention kernels +# are layout-agnostic; THD just adds slicing/concat glue), so any +# divergence beyond float-precision tolerance signals a plumbing bug. + + +def _make_packed_seq_params_thd(seg_lens, device='cuda'): + """Build a ``PackedSeqParams(qkv_format='thd', ...)`` from a list of + per-segment seq lengths. Self-attention contract: ``cu_seqlens_q == + cu_seqlens_kv``; ``*_padded`` mirrors the unpadded (no padding tested). + """ + cu_seqlens = torch.tensor( + [0] + list(torch.tensor(seg_lens, dtype=torch.int64).cumsum(0).tolist()), + dtype=torch.int32, + device=device, + ) + max_len = int(max(seg_lens)) if seg_lens else 0 + return PackedSeqParams( + cu_seqlens_q=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_kv_padded=cu_seqlens, + max_seqlen_q=max_len, + max_seqlen_kv=max_len, + qkv_format='thd', + ) + + +@pytest.mark.parametrize("compress_ratio", [4, 128]) +class TestCompressorThd: + """``Compressor`` THD-packed forward path + (``Compressor.forward(x, packed_seq_params=...)`` → ``_forward_thd``). + + Covers: + * Per-segment compressed-length contract: + ``cu_seqlens_compressed[b+1] - cu_seqlens_compressed[b] + == seqlen[b] // ratio``. + * Shape + dtype of the packed compressed-KV tensor. + * All-segments-too-short fast path (returns ``(None, cu_seqlens_compressed)``). + * B=1 single-segment THD matches SBHD-b=1 (numerical parity — same + per-segment math, just different layout glue). + * Gradient flow through the THD compression path. + """ + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config(csa_compress_ratios=[4, 128, 4, 128]) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + def _make_compressor(self, compress_ratio): + return Compressor( + config=self.config, + submodules=_make_compressor_submodules(), + compress_ratio=compress_ratio, + head_dim=self.config.v_head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_output_shape_and_cu_seqlens(self, compress_ratio): + """Multi-segment THD: each segment's compressed length is + ``seqlen[b] // ratio``; totals match the concat'd output. + """ + # Pick three segment lengths that each compress non-trivially. + seg_lens = [compress_ratio * 5, compress_ratio * 3, compress_ratio * 7] + total = sum(seg_lens) + packed = _make_packed_seq_params_thd(seg_lens) + compressor = self._make_compressor(compress_ratio) + + x = torch.randn(total, 1, self.config.hidden_size, dtype=torch.bfloat16, device='cuda') + out, cu_seqlens_compressed = compressor(x, packed_seq_params=packed) + + # Per-segment compressed lengths. + expected_per_seg = [s // compress_ratio for s in seg_lens] + expected_total = sum(expected_per_seg) + + assert out is not None + assert out.shape == (expected_total, 1, self.config.v_head_dim), ( + f"compressed_thd shape {tuple(out.shape)} != expected " + f"{(expected_total, 1, self.config.v_head_dim)}" + ) + assert out.dtype == torch.bfloat16 + # cu_seqlens_compressed[b+1] - cu_seqlens_compressed[b] == seqlen[b] // ratio. + diffs = (cu_seqlens_compressed[1:] - cu_seqlens_compressed[:-1]).cpu().tolist() + assert ( + diffs == expected_per_seg + ), f"cu_seqlens_compressed segment lengths {diffs} != {expected_per_seg}" + assert not torch.isnan(out).any() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_all_segments_too_short(self, compress_ratio): + """All segments shorter than ``ratio`` → returns + ``(None, cu_seqlens_compressed_all_zeros)``. + """ + seg_lens = [compress_ratio - 1, compress_ratio - 1] + total = sum(seg_lens) + packed = _make_packed_seq_params_thd(seg_lens) + compressor = self._make_compressor(compress_ratio) + + x = torch.randn(total, 1, self.config.hidden_size, dtype=torch.bfloat16, device='cuda') + out, cu_seqlens_compressed = compressor(x, packed_seq_params=packed) + + assert out is None + # All per-segment compressed lengths are zero. + diffs = (cu_seqlens_compressed[1:] - cu_seqlens_compressed[:-1]).cpu().tolist() + assert all( + d == 0 for d in diffs + ), f"all-short batch should have cu_seqlens_compressed all zero, got {diffs}" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_b1_matches_sbhd_b1(self, compress_ratio): + """B=1 single-segment THD output matches SBHD-b=1 on identical + input (compressor weights shared between the two calls). For the + same hidden states the per-segment math is identical, so the + outputs must agree within bf16-precision tolerance. + """ + seq_len = compress_ratio * 8 + compressor = self._make_compressor(compress_ratio) + + torch.manual_seed(42) + x_thd = torch.randn( + seq_len, 1, self.config.hidden_size, dtype=torch.bfloat16, device='cuda' + ) + # SBHD-b=1 input is the same data, no reshape needed (already (sq, 1, h)). + x_sbhd = x_thd + + # SBHD path: pass packed_seq_params=None → _forward_sbhd. + out_sbhd = compressor(x_sbhd, packed_seq_params=None) + + # THD path: pass packed_seq_params with single segment. + packed = _make_packed_seq_params_thd([seq_len]) + out_thd, cu_comp = compressor(x_thd, packed_seq_params=packed) + + assert out_sbhd is not None and out_thd is not None + assert ( + out_sbhd.shape == out_thd.shape + ), f"shape mismatch: sbhd={tuple(out_sbhd.shape)}, thd={tuple(out_thd.shape)}" + # cu_seqlens_compressed = [0, n_compressed]. + assert cu_comp[-1].item() == seq_len // compress_ratio + + # Numerical parity. bf16 + small per-segment-loop ordering differences + # mean we need a wider tol than fp32 would warrant. + assert torch.allclose(out_sbhd.float(), out_thd.float(), atol=5e-2, rtol=5e-2), ( + f"B=1 SBHD/THD parity failed: max abs diff = " + f"{(out_sbhd.float() - out_thd.float()).abs().max().item():.4e}" + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_gradient_flow(self, compress_ratio): + """Backward through Compressor THD populates grad on ``x`` and + every learnable parameter in the Compressor. + """ + seg_lens = [compress_ratio * 4, compress_ratio * 6] + total = sum(seg_lens) + packed = _make_packed_seq_params_thd(seg_lens) + compressor = self._make_compressor(compress_ratio) + + x = torch.randn( + total, 1, self.config.hidden_size, dtype=torch.bfloat16, device='cuda' + ).requires_grad_(True) + out, _ = compressor(x, packed_seq_params=packed) + loss = out.sum() + loss.backward() + + assert x.grad is not None and not torch.isnan(x.grad).any() + for name, p in compressor.named_parameters(): + if p.requires_grad: + assert p.grad is not None, f"Compressor param {name} has no grad" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_projection_disables_fp8(self, compress_ratio, monkeypatch): + """THD projections use the same high-precision path as SBHD projections.""" + compressor = self._make_compressor(compress_ratio) + tracker = _DisabledContextTracker() + calls = [] + + for name, projection in ( + ('linear_wkv', compressor.linear_wkv), + ('linear_wgate', compressor.linear_wgate), + ): + original_forward = projection.forward + + def checked_forward(*args, _name=name, _forward=original_forward, **kwargs): + assert tracker.depth > 0, f"{_name} ran outside the FP8-disabled context" + calls.append(_name) + return _forward(*args, **kwargs) + + monkeypatch.setattr(projection, 'forward', checked_forward) + + monkeypatch.setattr( + 'megatron.core.transformer.experimental_attention_variant.csa.get_fp8_disabled_context', + tracker, + ) + seq_len = compress_ratio * 2 + x = torch.randn(seq_len, 1, self.config.hidden_size, dtype=torch.bfloat16, device='cuda') + compressor(x, packed_seq_params=_make_packed_seq_params_thd([seq_len])) + + assert calls == ['linear_wkv', 'linear_wgate'] + assert tracker.entries == 1 + + +class TestCSAIndexerThd: + """``CSAIndexer`` THD-packed paths: + * ``forward_before_topk(packed_seq_params=...)`` — 4-tuple return + with ``cu_seqlens_compressed_idx``. + * ``forward(packed_seq_params=...)`` — THD dispatch through + :func:`fused_qk_topk_naive_thd`. New in the THD-completion turn. + + Multi-segment shape contract + B=1 SBHD-b=1 parity. + """ + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.compress_ratio = 4 + cls.config = _make_mla_config(csa_compress_ratios=[4, 4, 4, 4], dsa_indexer_topk=8) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + cls.indexer = CSAIndexer( + config=cls.config, + submodules=_make_csa_indexer_submodules(), + compress_ratio=cls.compress_ratio, + rotary_pos_emb=cls.rotary_pos_emb, + pg_collection=cls.pg_collection, + ).cuda() + + yield + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_forward_before_topk_returns_4_tuple(self): + """THD ``forward_before_topk`` returns + ``(q, k, weights, cu_seqlens_compressed_idx)`` with THD shapes + (dummy ``b=1`` dim retained for layout consistency with the SBHD + 4-D / 3-D contract that downstream THD callers ``.squeeze(1)``). + """ + ratio = self.compress_ratio + seg_lens = [ratio * 6, ratio * 4] + total = sum(seg_lens) + expected_total_comp = sum(s // ratio for s in seg_lens) + packed = _make_packed_seq_params_thd(seg_lens) + + x = torch.randn(total, 1, self.config.hidden_size, dtype=torch.bfloat16, device='cuda') + qr = torch.randn(total, 1, self.config.q_lora_rank, dtype=torch.bfloat16, device='cuda') + + result = self.indexer.forward_before_topk(x, qr, packed) + assert len(result) == 4, "THD forward_before_topk should return a 4-tuple" + q, k, weights, cu_seqlens_compressed_idx = result + + assert q.shape == ( + total, + 1, + self.config.dsa_indexer_n_heads, + self.config.dsa_indexer_head_dim, + ) + assert weights.shape == (total, 1, self.config.dsa_indexer_n_heads) + assert k.shape == (expected_total_comp, 1, self.config.dsa_indexer_head_dim) + # cu_seqlens_compressed_idx mirrors the compressor's cu_seqlens. + diffs = (cu_seqlens_compressed_idx[1:] - cu_seqlens_compressed_idx[:-1]).cpu().tolist() + assert diffs == [s // ratio for s in seg_lens] + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_forward_shape_and_dtype(self): + """THD ``forward`` returns ``(None, (total_q, topk) int64)``.""" + ratio = self.compress_ratio + seg_lens = [ratio * 5, ratio * 3] + total = sum(seg_lens) + packed = _make_packed_seq_params_thd(seg_lens) + + x = torch.randn(total, 1, self.config.hidden_size, dtype=torch.bfloat16, device='cuda') + qr = torch.randn(total, 1, self.config.q_lora_rank, dtype=torch.bfloat16, device='cuda') + + index_scores, topk = self.indexer(x, qr, packed_seq_params=packed) + # THD return contract: per-segment scores aren't surfaced + # (heterogeneous shapes); only consumers in csa.py + # force_unfused inference use this path and discard scores. + assert index_scores is None + assert topk.shape == (total, self.config.dsa_indexer_topk) + assert topk.dtype == torch.int64 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_per_segment_kv_scope(self): + """Top-K LOCAL ids stay in ``[0, seqlen_compressed[b])`` per-segment + (NOT flat-global ids into the concat'd indexer-K). + """ + ratio = self.compress_ratio + seg_lens = [ratio * 8, ratio * 4] + total = sum(seg_lens) + n_comp_per_seg = [s // ratio for s in seg_lens] + packed = _make_packed_seq_params_thd(seg_lens) + + x = torch.randn(total, 1, self.config.hidden_size, dtype=torch.bfloat16, device='cuda') + qr = torch.randn(total, 1, self.config.q_lora_rank, dtype=torch.bfloat16, device='cuda') + + _, topk = self.indexer(x, qr, packed_seq_params=packed) + + # Segment 0 rows: valid ids must be in [0, n_comp_per_seg[0]). + seg0 = topk[: seg_lens[0]] + seg0_valid = seg0[seg0 >= 0] + if seg0_valid.numel() > 0: + assert (seg0_valid < n_comp_per_seg[0]).all(), ( + f"segment 0 ids out of range: max={seg0_valid.max().item()}, " + f"expected < {n_comp_per_seg[0]}" + ) + # Segment 1 rows: valid ids must be in [0, n_comp_per_seg[1]). + seg1 = topk[seg_lens[0] :] + seg1_valid = seg1[seg1 >= 0] + if seg1_valid.numel() > 0: + assert (seg1_valid < n_comp_per_seg[1]).all(), ( + f"segment 1 ids out of range: max={seg1_valid.max().item()}, " + f"expected < {n_comp_per_seg[1]}" + ) + + +class TestCompressedSparseAttentionThd: + """End-to-end ``CompressedSparseAttention(packed_seq_params=...)`` + integration tests covering all THD-supported Path × fused/force_unfused + combinations. Each test verifies no NaN + expected output shape; the + deep numerical correctness is established at lower layers by the + real-kernel parity tests (``TestRealKernelFusedIndexerSparseAttn*``, + ``TestFusedDSAIndexerLossThd``, ``TestFusedQkTopkNaiveThd``). + + THD output shape is ``(total_q, 1, np * v_head_dim)`` — the dummy + ``b=1`` axis is re-added inside ``_forward_thd`` so downstream + callers can keep the SBHD ``(seq, batch, hidden)`` 3-D contract. + """ + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config( + csa_compress_ratios=[4, 128, 4, 128], + csa_window_size=8, + dsa_indexer_topk=8, + dsa_indexer_loss_coeff=1.0, + ) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + def _get_layer_number(self, compress_ratio): + """Return the (1-indexed) layer number whose + ``csa_compress_ratios`` entry matches ``compress_ratio``.""" + for i, r in enumerate(self.config.csa_compress_ratios): + if r == compress_ratio: + return i + 1 + raise ValueError(f"No layer with compress_ratio={compress_ratio}") + + def _build_csa(self, compress_ratio, *, force_unfused_dsa=False): + # ``force_unfused_dsa`` is a config-level attribute consumed by + # ``CompressedSparseAttention.__init__`` via ``getattr(config, + # 'force_unfused_dsa', False)``; set it on the config object + # before constructing the module. + self.config.force_unfused_dsa = force_unfused_dsa + return CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=self._get_layer_number(compress_ratio), + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, + ).cuda() + + def _make_thd_inputs(self, seg_lens): + """Build a ``(query, key, value, x, qr, packed_seq_params)`` + tuple for a multi-segment THD batch of given segment lengths. + """ + total = sum(seg_lens) + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + query = torch.randn(total, np_, hn, dtype=torch.bfloat16, device='cuda') + key = torch.randn(total, 1, 1, hn, dtype=torch.bfloat16, device='cuda') + value = key.clone() + x = torch.randn(total, 1, self.config.hidden_size, dtype=torch.bfloat16, device='cuda') + qr = torch.randn(total, 1, self.config.q_lora_rank, dtype=torch.bfloat16, device='cuda') + packed = _make_packed_seq_params_thd(seg_lens) + return query, key, value, x, qr, packed + + # ---- Path A (compress_ratio=128: indexer disabled, all-compressed) ---- + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_path_a_forward(self): + """Path A (THD): compress_ratio=128 → indexer=None → attend to + ALL compressed positions per segment via ``get_compress_topk_idxs_thd``. + """ + compress_ratio = 128 + csa = self._build_csa(compress_ratio) + # Make segment lengths long enough that each compresses ≥1 position. + seg_lens = [compress_ratio * 2 + 50, compress_ratio + 30] + total = sum(seg_lens) + query, key, value, x, qr, packed = self._make_thd_inputs(seg_lens) + + csa.eval() + with torch.no_grad(): + output = csa( + query=query, + key=key, + value=value, + attention_mask=None, + x=x, + qr=qr, + packed_seq_params=packed, + ) + np_ = self.config.num_attention_heads + assert output.shape == (total, 1, np_ * self.config.v_head_dim) + assert not torch.isnan(output).any() + + # ---- Path B (compress_ratio=4, training): fused × sparse/dense × force_unfused ---- + + @pytest.mark.parametrize( + "sparse_loss, force_unfused_dsa", + [ + (False, False), # fused, dense loss (cuDNN dense kernels) + (True, False), # fused, sparse loss (cuDNN sparse kernels) + (True, True), # force_unfused (PyTorch ref) — uses + # config.dsa_indexer_use_sparse_loss directly + ], + ids=['fused_dense', 'fused_sparse', 'force_unfused'], + ) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_path_b_training_forward_backward(self, sparse_loss, force_unfused_dsa): + """Path B (THD training): all three supported combos exercise + the indexer + KL-loss path with grad flow through Q/K/x/qr. + """ + # Set the sparse-loss config flag (read inside _forward_thd). + self.config.dsa_indexer_use_sparse_loss = sparse_loss + + compress_ratio = 4 + csa = self._build_csa(compress_ratio, force_unfused_dsa=force_unfused_dsa) + # Multi-segment with enough length for indexer top-K to be exercised. + seg_lens = [compress_ratio * 16, compress_ratio * 8] + total = sum(seg_lens) + query, key, value, x, qr, packed = self._make_thd_inputs(seg_lens) + + # Require grad on differentiable inputs (mirrors the SBHD backward test). + query.requires_grad_(True) + key.requires_grad_(True) + x.requires_grad_(True) + qr.requires_grad_(True) + + csa.train() + output = csa( + query=query, + key=key, + value=value, + attention_mask=None, + x=x, + qr=qr, + packed_seq_params=packed, + ) + np_ = self.config.num_attention_heads + assert output.shape == (total, 1, np_ * self.config.v_head_dim) + assert not torch.isnan(output).any() + + # Backward: indexer loss is attached via DSAIndexerLossAutoScaler so + # ``output.sum().backward()`` triggers grads through both the attn + # output path AND the indexer-loss path. + output.sum().backward() + # Differentiable leaves should have grads. + assert query.grad is not None and not torch.isnan(query.grad).any() + assert key.grad is not None and not torch.isnan(key.grad).any() + # CSA params (compressor + indexer + attn_sink) should be reached. + seen_any_param_grad = False + for name, p in csa.named_parameters(): + if p.requires_grad and p.grad is not None: + seen_any_param_grad = True + assert not torch.isnan(p.grad).any(), f"param {name} grad has NaN" + assert seen_any_param_grad, "no CSA param received a gradient" + + # ---- Path C (compress_ratio=4, inference): fused × force_unfused ---- + + @pytest.mark.parametrize("force_unfused_dsa", [False, True], ids=['fused', 'force_unfused']) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_path_c_inference_forward(self, force_unfused_dsa): + """Path C (THD inference): indexer top-K + sparse attn, no loss. + Both the cuDNN fused path and the PyTorch-ref force_unfused path + produce a well-formed output. + """ + compress_ratio = 4 + csa = self._build_csa(compress_ratio, force_unfused_dsa=force_unfused_dsa) + seg_lens = [compress_ratio * 16, compress_ratio * 12] + total = sum(seg_lens) + query, key, value, x, qr, packed = self._make_thd_inputs(seg_lens) + + csa.eval() + with torch.no_grad(): + output = csa( + query=query, + key=key, + value=value, + attention_mask=None, + x=x, + qr=qr, + packed_seq_params=packed, + ) + np_ = self.config.num_attention_heads + assert output.shape == (total, 1, np_ * self.config.v_head_dim) + assert not torch.isnan(output).any() + + # ---- B=1 SBHD/THD parity (one happy-path sanity check) ---- + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_b1_sbhd_thd_parity_inference_path_c(self): + """B=1 single-segment THD inference output matches SBHD-b=1 on + the same data (force_unfused path → fully deterministic, no + cuDNN/FlashMLA topk-tie nondeterminism). + + Wider tol than the kernel-level tests because the full CSA + forward chains many bf16 ops together; we just verify "no + plumbing bug" rather than tight numerical equality. + """ + compress_ratio = 4 + # force_unfused → uses the PyTorch indexer reference (no cuDNN + # radix-topK tie-breaking nondeterminism). + csa = self._build_csa(compress_ratio, force_unfused_dsa=True) + sq = compress_ratio * 16 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + torch.manual_seed(7) + query = torch.randn(sq, 1, np_, hn, dtype=torch.bfloat16, device='cuda') + key = torch.randn(sq, 1, 1, hn, dtype=torch.bfloat16, device='cuda') + value = key.clone() + x = torch.randn(sq, 1, self.config.hidden_size, dtype=torch.bfloat16, device='cuda') + qr = torch.randn(sq, 1, self.config.q_lora_rank, dtype=torch.bfloat16, device='cuda') + + csa.eval() + with torch.no_grad(): + # SBHD path: packed_seq_params=None. + out_sbhd = csa( + query=query, + key=key, + value=value, + attention_mask=None, + x=x, + qr=qr, + packed_seq_params=None, + ) + # THD path: single-segment packed_seq_params. Query is 3-D + # ``(total_q, np, hn)`` per TE THD convention, so drop the + # SBHD b=1 head dimension for the THD call. + packed = _make_packed_seq_params_thd([sq]) + out_thd = csa( + query=query.squeeze(1), + key=key, + value=value, + attention_mask=None, + x=x, + qr=qr, + packed_seq_params=packed, + ) + + # SBHD output: (sq, 1, np*hn). THD output: (sq, 1, np*hn). Same shape. + assert out_sbhd.shape == out_thd.shape + assert torch.allclose(out_sbhd.float(), out_thd.float(), atol=5e-2, rtol=5e-2), ( + f"SBHD/THD B=1 parity failed: max abs diff = " + f"{(out_sbhd.float() - out_thd.float()).abs().max().item():.4e}" + ) + + +# =========================================================================== +# _apply_rope direct THD tests (4 corners: ratio={1, >1} × fused={False, True}) +# =========================================================================== +# +# Direct tests of ``_apply_rope`` THD branches. Previously these were +# only exercised indirectly via ``Compressor._forward_thd`` (ratio>1) +# and ``CSAIndexer.forward_before_topk`` (ratio=1). Direct tests give +# clearer failure attribution and pin down the contract for each of the +# 4 supported (ratio, apply_rope_fusion) combinations. + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +class TestApplyRopeThd: + """Direct tests of :func:`_apply_rope` THD branches. + + The class is parametrized over ``rope_type`` (``"rope"`` / + ``"yarn"``), so every test runs with both ``RotaryEmbedding`` + and ``YarnRotaryEmbedding``. + + For each rope type the function has four distinct THD paths: + * (ratio=1, fused=False): forward ``cu_seqlens`` to the rotary + module's packed mode + ``apply_rotary_pos_emb``. + * (ratio=1, fused=True): forward ``cu_seqlens`` to the fused MLA + RoPE kernel. + * (ratio>1, fused=False): build a per-segment-strided rotary + table by slicing a global ``max_seg * ratio`` table with stride + ``ratio`` per segment, concat into a packed table aligned with + ``cu_seqlens``, then ``apply_rotary_pos_emb``. + * (ratio>1, fused=True): same per-segment-strided slice + concat + construction but applied to cos/sin tables instead of the + rotary embedding tensor, fed to the fused kernel. + + For each corner we verify: + * Output shape == input shape (RoPE is in-place w.r.t. shape). + * No NaN in the output (per-segment). + * B=1 single-segment THD output matches the equivalent SBHD-b=1 + call on the same input (numerical parity within bf16 tol). + """ + + @pytest.fixture(scope='class', autouse=True, params=["rope", "yarn"], ids=["rope", "yarn"]) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + rope_type = request.param + cls = request.cls + cls.rope_type = rope_type + cls.config = _make_mla_config(rope_type=rope_type) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + if rope_type == "yarn": + from megatron.core.models.common.embeddings import YarnRotaryEmbedding + + cls.rotary_pos_emb = YarnRotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_base=cls.config.rotary_base, + scaling_factor=cls.config.rotary_scaling_factor, + original_max_position_embeddings=cls.config.original_max_position_embeddings, + beta_fast=cls.config.beta_fast, + beta_slow=cls.config.beta_slow, + mscale=cls.config.mscale, + mscale_all_dim=cls.config.mscale_all_dim, + cp_group=cls.pg_collection.cp, + ) + else: + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + cls.pos_dim = cls.config.qk_pos_emb_head_dim + cls.nope_dim = cls.config.v_head_dim - cls.pos_dim + cls.head_dim = cls.config.v_head_dim + + yield + Utils.destroy_model_parallel() + + def _make_input_thd(self, total_q): + # 3-D input ``(seq, batch=1, head_dim)`` — the shape that + # ``Compressor._forward_thd`` and ``CSAIndexer.forward_before_topk`` + # feed in (with the dummy ``b=1`` axis preserved). ``_apply_rope`` + # also accepts 4-D (with explicit head dim); both branches go + # through the same code path after a temporary head-dim insert. + return torch.randn(total_q, 1, self.head_dim, dtype=torch.bfloat16, device='cuda') + + @pytest.mark.parametrize("ratio", [1, 4], ids=["ratio_1", "ratio_4"]) + @pytest.mark.parametrize("apply_rope_fusion", [False, True], ids=["unfused", "fused"]) + def test_thd_shape_and_no_nan(self, ratio, apply_rope_fusion): + """All 4 corners produce same-shape, NaN-free output for a + multi-segment THD batch. + """ + prev_fusion = self.config.apply_rope_fusion + self.config.apply_rope_fusion = apply_rope_fusion + try: + seg_lens = [16, 24, 8] + total = sum(seg_lens) + x = self._make_input_thd(total) + cu_seqlens = _cu_seqlens(seg_lens, device='cuda') + + out = _apply_rope( + x, + self.nope_dim, + self.pos_dim, + self.rotary_pos_emb, + self.config, + rotary_seq_len=0, # unused when cu_seqlens supplied + ratio=ratio, + cp_group=self.pg_collection.cp, + cu_seqlens=cu_seqlens, + max_seqlen_rope=max(seg_lens) * ratio, + ) + + tag = f"(rope={self.rope_type}, ratio={ratio}, fused={apply_rope_fusion})" + assert ( + out.shape == x.shape + ), f"{tag}: shape {tuple(out.shape)} != input {tuple(x.shape)}" + offset = 0 + for i, seg_len in enumerate(seg_lens): + assert not torch.isnan( + out[offset : offset + seg_len] + ).any(), f"{tag}: segment {i} produced NaN" + offset += seg_len + finally: + self.config.apply_rope_fusion = prev_fusion + + @pytest.mark.parametrize("ratio", [1, 4], ids=["ratio_1", "ratio_4"]) + @pytest.mark.parametrize("apply_rope_fusion", [False, True], ids=["unfused", "fused"]) + def test_thd_b1_matches_sbhd_b1(self, ratio, apply_rope_fusion): + """B=1 single-segment THD matches SBHD-b=1 on the same input + for all 4 corners. The two paths build their rotary tables + independently but for a single segment with ``cu_seqlens = [0, + sq]`` they should produce numerically identical output. + """ + prev_fusion = self.config.apply_rope_fusion + self.config.apply_rope_fusion = apply_rope_fusion + try: + sq = 16 + x = self._make_input_thd(sq) + cu_seqlens = _cu_seqlens([sq], device='cuda') + + # SBHD: cu_seqlens=None. For ratio>1 the SBHD branch slices + # a length ``sq*ratio`` table with stride ratio. ``x`` must + # have a sequence-first layout (which it does: (sq, 1, head_dim)). + out_sbhd = _apply_rope( + x.clone(), + self.nope_dim, + self.pos_dim, + self.rotary_pos_emb, + self.config, + rotary_seq_len=sq, + ratio=ratio, + cp_group=self.pg_collection.cp, + cu_seqlens=None, + ) + out_thd = _apply_rope( + x.clone(), + self.nope_dim, + self.pos_dim, + self.rotary_pos_emb, + self.config, + rotary_seq_len=0, # unused for THD + ratio=ratio, + cp_group=self.pg_collection.cp, + cu_seqlens=cu_seqlens, + max_seqlen_rope=sq * ratio, + ) + + tag = f"(rope={self.rope_type}, ratio={ratio}, fused={apply_rope_fusion})" + assert out_sbhd.shape == out_thd.shape, f"{tag} shape mismatch" + assert torch.allclose(out_sbhd.float(), out_thd.float(), atol=1e-2, rtol=1e-2), ( + f"{tag} SBHD/THD B=1 parity failed: max abs diff = " + f"{(out_sbhd.float() - out_thd.float()).abs().max().item():.4e}" + ) + finally: + self.config.apply_rope_fusion = prev_fusion + + +class TestCSAHighPrecisionParams: + """The compressor ``ape`` and attention ``attn_sink`` parameters must stay in FP32 + (the reference DeepSeek V4 checkpoint stores them in FP32), even after the model is + converted to BF16/FP16 by ``Float16Module``.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config(csa_compress_ratios=[4, 4, 4, 4]) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_ape_and_attn_sink_stay_fp32_after_bf16_conversion(self): + from megatron.core.transformer.module import Float16Module + + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=4, + name="decoder.layers.0.self_attention.core_attention", + ) + + assert csa.attn_sink.dtype == torch.float32 + assert csa.compressor.ape.dtype == torch.float32 + assert csa.indexer.compressor.ape.dtype == torch.float32 + + bf16_module = Float16Module(config=self.config, module=csa) + + assert bf16_module.module.attn_sink.dtype == torch.float32 + assert bf16_module.module.compressor.ape.dtype == torch.float32 + assert bf16_module.module.indexer.compressor.ape.dtype == torch.float32 + assert bf16_module.module.compressor.linear_wkv.weight.dtype == torch.bfloat16 + assert bf16_module.module.compressor.linear_wgate.weight.dtype == torch.bfloat16 diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py index 135c4802dd3..b8e999f3d70 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py @@ -32,6 +32,7 @@ _validate_nonpacked_cp_uniform_length, compute_dsa_indexer_loss, fused_qk_topk_naive, + fused_qk_topk_naive_thd, is_dsa_skip_topk_layer, rotate_activation, source_dsa_compute_layer, @@ -3898,3 +3899,446 @@ def test_get_dsa_module_spec_rejects_qk_l2_norm(self): config = self._make_dsa_config(qk_l2_norm=True) with pytest.raises(AssertionError, match="qk_l2_norm is not supported"): get_dsa_module_spec_for_backend(config, backend=None) + + +# =========================================================================== +# THD: FusedDSAIndexerLoss +# =========================================================================== + + +class TestFusedDSAIndexerLossThd: + """``FusedDSAIndexerLoss`` THD branch — per-segment loop that delegates + each segment to the SBHD naive helpers with ``b=1`` and aggregates + via row-weighted-mean. + + For a single-segment THD batch (``cu_seqlens_q = [0, sq]``) the THD + invocation must produce numerically equivalent loss + gradients as + the SBHD invocation with ``b=1`` on the same data — the only + difference is the (B=1) per-segment slicing/concat glue. + """ + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + request.cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['tp'] + ) + yield + Utils.destroy_model_parallel() + + @pytest.mark.parametrize('sparse_loss', [False, True], ids=['dense_loss', 'sparse_loss']) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_single_segment_matches_sbhd_b1(self, sparse_loss): + """B=1 THD invocation should match the equivalent SBHD-b=1 call + (loss + gradients) for both sparse and dense KL loss variants. + """ + torch.manual_seed(0) + sq = 32 + n_compressed = sq // 4 # ratio=4 → compressed K len per segment + ratio = 4 + num_heads = 4 + head_dim = 64 + idx_nh, idx_hd = 4, 32 + topk = 4 + softmax_scale = head_dim**-0.5 + loss_coeff = 0.5 + + # ---- Common inputs (SBHD-shape with b=1) ----------------------- + def _rand(*shape, dtype=torch.float32): + return torch.randn(*shape, dtype=dtype, device='cuda') + + q_sbhd = _rand(sq, 1, idx_nh, idx_hd).requires_grad_(True) + w_sbhd = _rand(sq, 1, idx_nh).requires_grad_(True) + k_sbhd = _rand(n_compressed, 1, idx_hd).requires_grad_(True) + query_sbhd = _rand(sq, 1, num_heads, head_dim, dtype=torch.bfloat16) + key_sbhd = _rand(n_compressed, 1, num_heads, head_dim, dtype=torch.bfloat16) + + # SBHD per-batch causal mask: (1, sq, n_compressed). + cols = torch.arange(n_compressed, device='cuda').unsqueeze(0).expand(sq, -1) + positions = torch.arange(1, sq + 1, device='cuda').unsqueeze(1) + mask_sbhd = torch.where(cols >= positions // ratio, float('-inf'), 0.0).unsqueeze(0) + + # ---- SBHD reference -------------------------------------------- + topk_indices_sbhd, loss_sbhd = FusedDSAIndexerLoss.apply( + q_sbhd, + w_sbhd, + k_sbhd, + query_sbhd, + key_sbhd, + softmax_scale, + topk, + loss_coeff, + mask_sbhd, + sparse_loss, + self.pg_collection, + None, # varlen_starts + None, # varlen_ends + None, # key_positions + None, # query_valid_rows + False, # calculate_per_token_loss + ) + loss_sbhd.backward() + grad_q_sbhd = q_sbhd.grad.clone() + grad_w_sbhd = w_sbhd.grad.clone() + grad_k_sbhd = k_sbhd.grad.clone() + + # ---- THD equivalent (B=1, total_q=sq) -------------------------- + q_thd = q_sbhd.detach().squeeze(1).clone().requires_grad_(True) + w_thd = w_sbhd.detach().squeeze(1).clone().requires_grad_(True) + k_thd = k_sbhd.detach().squeeze(1).clone().requires_grad_(True) + query_thd = query_sbhd.squeeze(1) + key_thd = key_sbhd.squeeze(1) + + cu_seqlens_q = torch.tensor([0, sq], dtype=torch.int32, device='cuda') + cu_seqlens_comp = torch.tensor([0, n_compressed], dtype=torch.int32, device='cuda') + + topk_indices_thd, loss_thd = FusedDSAIndexerLoss.apply( + q_thd, + w_thd, + k_thd, + query_thd, + key_thd, + softmax_scale, + topk, + loss_coeff, + None, # mask: built per-segment internally for THD + sparse_loss, + self.pg_collection, + None, # varlen_starts + None, # varlen_ends + None, # key_positions + None, # query_valid_rows + False, # calculate_per_token_loss + True, # use_relu + cu_seqlens_q, + cu_seqlens_comp, + ratio, + ) + loss_thd.backward() + grad_q_thd = q_thd.grad + grad_w_thd = w_thd.grad + grad_k_thd = k_thd.grad + + tag = f"[sparse={sparse_loss}]" + + # Loss + grads must match the SBHD-b=1 reference (same math, + # same data; only the slicing-and-concat glue differs). + assert torch.allclose(loss_thd, loss_sbhd, rtol=1e-5, atol=1e-5), ( + f"{tag} loss mismatch: thd={loss_thd.item()}, " f"sbhd={loss_sbhd.item()}" + ) + # topk_indices_thd is (total_q, topk); SBHD is (1, sq, topk). + assert torch.equal( + topk_indices_thd, topk_indices_sbhd.squeeze(0).int() + ), f"{tag} topk mismatch" + assert torch.allclose( + grad_q_thd, grad_q_sbhd.squeeze(1), rtol=1e-5, atol=1e-5 + ), f"{tag} grad_q mismatch" + assert torch.allclose( + grad_w_thd, grad_w_sbhd.squeeze(1), rtol=1e-5, atol=1e-5 + ), f"{tag} grad_w mismatch" + assert torch.allclose( + grad_k_thd, grad_k_sbhd.squeeze(1), rtol=1e-5, atol=1e-5 + ), f"{tag} grad_k mismatch" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_missing_kwarg_raises(self): + """THD mode requires both ``cu_seqlens_compressed_idx`` and + ``ratio``; supplying ``cu_seqlens_q`` alone raises ``ValueError``. + """ + sq, n_compressed = 8, 2 + idx_nh, idx_hd = 4, 32 + num_heads, head_dim = 4, 64 + q = torch.zeros(sq, idx_nh, idx_hd, dtype=torch.float32, device='cuda') + w = torch.zeros(sq, idx_nh, dtype=torch.float32, device='cuda') + k = torch.zeros(n_compressed, idx_hd, dtype=torch.float32, device='cuda') + query = torch.zeros(sq, num_heads, head_dim, dtype=torch.bfloat16, device='cuda') + key = torch.zeros(n_compressed, num_heads, head_dim, dtype=torch.bfloat16, device='cuda') + cu_q = torch.tensor([0, sq], dtype=torch.int32, device='cuda') + with pytest.raises(ValueError, match="THD mode requires"): + FusedDSAIndexerLoss.apply( + q, + w, + k, + query, + key, + head_dim**-0.5, + 2, + 1.0, + None, + False, + self.pg_collection, + None, # varlen_starts + None, # varlen_ends + None, # key_positions + None, # query_valid_rows + False, # calculate_per_token_loss + True, # use_relu + cu_q, # cu_seqlens_q supplied + None, # cu_seqlens_compressed_idx MISSING + None, # ratio MISSING + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_multiseg_zero_compressed_segment_row_mean_normalization(self): + """THD aggregated loss should be row-mean over ``total_q`` even when + some segments have ``seqlen_compressed == 0``. + + Build a two-segment THD batch where segment 0 has no compressed keys + (sq=2, ratio=4 -> 0 compressed) and segment 1 has compressed keys + (sq=8 -> 2 compressed). Compare THD loss to an SBHD per-segment + reference aggregated as ``sum(loss_b * sq_b) / total_q``. + """ + torch.manual_seed(7) + seg_q_lens = [2, 8] + seg_comp_lens = [0, 2] + total_q = sum(seg_q_lens) + total_comp = sum(seg_comp_lens) + ratio = 4 + idx_nh, idx_hd = 4, 32 + num_heads, head_dim = 4, 64 + topk = 2 + softmax_scale = head_dim**-0.5 + loss_coeff = 0.5 + dev = 'cuda' + + # THD inputs + q_thd = torch.randn(total_q, idx_nh, idx_hd, dtype=torch.float32, device=dev) + w_thd = torch.randn(total_q, idx_nh, dtype=torch.float32, device=dev) + k_thd = torch.randn(total_comp, idx_hd, dtype=torch.float32, device=dev) + query_thd = torch.randn(total_q, num_heads, head_dim, dtype=torch.bfloat16, device=dev) + key_thd = torch.randn(total_comp, num_heads, head_dim, dtype=torch.bfloat16, device=dev) + + cu_seqlens_q = torch.tensor([0, 2, 10], dtype=torch.int32, device=dev) + cu_seqlens_comp = torch.tensor([0, 0, 2], dtype=torch.int32, device=dev) + + _, loss_thd = FusedDSAIndexerLoss.apply( + q_thd, + w_thd, + k_thd, + query_thd, + key_thd, + softmax_scale, + topk, + loss_coeff, + None, + False, + self.pg_collection, + None, # varlen_starts + None, # varlen_ends + None, # key_positions + None, # query_valid_rows + False, # calculate_per_token_loss + True, # use_relu + cu_seqlens_q, + cu_seqlens_comp, + ratio, + ) + + # SBHD per-segment reference: segment 0 contributes zero because it has + # no compressed keys; segment 1 contributes normally. + weighted_losses = [] + for b, (sq_b, sk_b) in enumerate(zip(seg_q_lens, seg_comp_lens)): + if sk_b == 0: + continue + q_start = int(cu_seqlens_q[b].item()) + q_end = int(cu_seqlens_q[b + 1].item()) + k_start = int(cu_seqlens_comp[b].item()) + k_end = int(cu_seqlens_comp[b + 1].item()) + + q_b = q_thd[q_start:q_end].unsqueeze(1) + w_b = w_thd[q_start:q_end].unsqueeze(1) + k_b = k_thd[k_start:k_end].unsqueeze(1) + query_b = query_thd[q_start:q_end].unsqueeze(1) + key_b = key_thd[k_start:k_end].unsqueeze(1) + + cols = torch.arange(sk_b, device=dev).unsqueeze(0).expand(sq_b, -1) + positions = torch.arange(1, sq_b + 1, device=dev).unsqueeze(1) + mask_b = torch.where(cols >= positions // ratio, float('-inf'), 0.0).unsqueeze(0) + + _, loss_b = FusedDSAIndexerLoss.apply( + q_b, + w_b, + k_b, + query_b, + key_b, + softmax_scale, + topk, + loss_coeff, + mask_b, + False, + self.pg_collection, + None, # varlen_starts + None, # varlen_ends + None, # key_positions + None, # query_valid_rows + False, # calculate_per_token_loss + ) + weighted_losses.append(loss_b * sq_b) + + expected_loss = torch.stack(weighted_losses).sum() / float(total_q) + assert torch.allclose(loss_thd, expected_loss, rtol=1e-5, atol=1e-5), ( + f"THD loss should be row-mean over total_q={total_q}: " + f"thd={loss_thd.item()}, expected={expected_loss.item()}" + ) + + +# =========================================================================== +# THD: fused_qk_topk_naive_thd (force_unfused_dsa + indexer + inference path) +# =========================================================================== + + +class TestFusedQkTopkNaiveThd: + """``fused_qk_topk_naive_thd`` — per-segment naive PyTorch QK + top-K + used by the THD ``force_unfused_dsa + indexer + inference`` path + (i.e., the THD branch of :meth:`CSAIndexer.forward`). + + Coverage: + * B=1 single-segment THD matches SBHD-b=1 ``fused_qk_topk_naive`` + ranking (same scores → same top-K positions among valid rows). + * Output shape + dtype contract. + * ``-1`` sentinel marking on invalid tail positions (rows whose + causal-valid count is < topk). + * Multi-segment dispatch isolates per-segment KV scopes (segment + ``b``'s top-K can only reference KV positions in + ``[0, seqlen_kv[b])``). + """ + + def _build_causal_mask(self, sq, sk, ratio, device): + """SBHD-shape ``(1, sq, sk)`` causal mask (mirrors + ``_build_causal_mask_seg`` for the reference path).""" + cols = torch.arange(sk, device=device).unsqueeze(0).expand(sq, -1) + positions = torch.arange(1, sq + 1, device=device).unsqueeze(1) + return torch.where(cols >= positions // ratio, float('-inf'), 0.0).unsqueeze(0) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_single_segment_matches_sbhd(self): + """B=1 single-segment THD top-K should match the SBHD-b=1 + reference among valid rows (rows whose causal-valid count is + smaller than ``topk`` get ``-1`` sentinels in THD where SBHD + returns garbage tail; we compare only the valid prefix). + """ + torch.manual_seed(0) + sq, n_compressed = 32, 8 + idx_nh, idx_hd = 4, 32 + topk = 4 + ratio = 4 + dev = 'cuda' + + q_thd = torch.randn(sq, idx_nh, idx_hd, dtype=torch.float32, device=dev) + k_thd = torch.randn(n_compressed, idx_hd, dtype=torch.float32, device=dev) + w_thd = torch.randn(sq, idx_nh, dtype=torch.float32, device=dev) + + cu_q = torch.tensor([0, sq], dtype=torch.int32, device=dev) + cu_kv = torch.tensor([0, n_compressed], dtype=torch.int32, device=dev) + + _, topk_thd = fused_qk_topk_naive_thd(q_thd, k_thd, w_thd, topk, cu_q, cu_kv, ratio) + + # SBHD reference: same data with b=1 + caller-supplied mask. + q_sbhd = q_thd.unsqueeze(1) + k_sbhd = k_thd.unsqueeze(1) + w_sbhd = w_thd.unsqueeze(1) + mask_sbhd = self._build_causal_mask(sq, n_compressed, ratio, dev) + _, topk_sbhd = fused_qk_topk_naive(q_sbhd, k_sbhd, w_sbhd, topk, mask_sbhd) + topk_sbhd = topk_sbhd.squeeze(0) # (sq, topk) + + # Per-row: compare only the leading ``n_valid`` slots; THD marks + # the rest as -1, SBHD's tail is undefined (masked -inf + # positions, ties may break differently). + for row in range(sq): + n_valid = min((row + 1) // ratio, n_compressed, topk) + assert torch.equal( + topk_thd[row, :n_valid].cpu(), topk_sbhd[row, :n_valid].cpu() + ), f"row {row}: top-K mismatch among valid slots" + assert ( + topk_thd[row, n_valid:] == -1 + ).all(), f"row {row}: THD must mark invalid tail as -1" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_shape_and_dtype(self): + """Returns ``(None, (total_q, topk) int64)``.""" + torch.manual_seed(0) + sq_a, sq_b = 8, 4 + kv_a, kv_b = 2, 1 + idx_nh, idx_hd = 2, 16 + topk = 3 + ratio = 4 + dev = 'cuda' + + q = torch.randn(sq_a + sq_b, idx_nh, idx_hd, dtype=torch.float32, device=dev) + k = torch.randn(kv_a + kv_b, idx_hd, dtype=torch.float32, device=dev) + w = torch.randn(sq_a + sq_b, idx_nh, dtype=torch.float32, device=dev) + cu_q = torch.tensor([0, sq_a, sq_a + sq_b], dtype=torch.int32, device=dev) + cu_kv = torch.tensor([0, kv_a, kv_a + kv_b], dtype=torch.int32, device=dev) + + scores, topk_idxs = fused_qk_topk_naive_thd(q, k, w, topk, cu_q, cu_kv, ratio) + assert scores is None + assert topk_idxs.shape == (sq_a + sq_b, topk) + assert topk_idxs.dtype == torch.int64 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_per_segment_kv_scope(self): + """Segment ``b``'s top-K LOCAL ids must live in + ``[0, seqlen_kv[b])`` (per-segment scope) — they are NOT + flat-global ids into the concatenated K tensor. + """ + torch.manual_seed(0) + sq_a, sq_b = 16, 12 + kv_a, kv_b = 4, 3 + idx_nh, idx_hd = 2, 16 + topk = 2 + ratio = 4 + dev = 'cuda' + + q = torch.randn(sq_a + sq_b, idx_nh, idx_hd, dtype=torch.float32, device=dev) + k = torch.randn(kv_a + kv_b, idx_hd, dtype=torch.float32, device=dev) + w = torch.randn(sq_a + sq_b, idx_nh, dtype=torch.float32, device=dev) + cu_q = torch.tensor([0, sq_a, sq_a + sq_b], dtype=torch.int32, device=dev) + cu_kv = torch.tensor([0, kv_a, kv_a + kv_b], dtype=torch.int32, device=dev) + + _, topk_idxs = fused_qk_topk_naive_thd(q, k, w, topk, cu_q, cu_kv, ratio) + + # Segment 0 rows: valid ids must be in [0, kv_a). + seg0 = topk_idxs[:sq_a] + seg0_valid = seg0[seg0 >= 0] + assert (seg0_valid < kv_a).all(), ( + f"segment 0 has out-of-range ids: max = {seg0_valid.max().item()}, " + f"expected < {kv_a}" + ) + # Segment 1 rows: valid ids must be in [0, kv_b). + seg1 = topk_idxs[sq_a:] + seg1_valid = seg1[seg1 >= 0] + assert (seg1_valid < kv_b).all(), ( + f"segment 1 has out-of-range ids: max = {seg1_valid.max().item()}, " + f"expected < {kv_b}" + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_invalid_tail_marked_minus_one(self): + """Early rows where ``(pos+1)//ratio < topk`` should have ``-1`` + sentinels in the tail of their top-K row. + """ + torch.manual_seed(0) + sq, n_compressed = 8, 4 + idx_nh, idx_hd = 2, 16 + topk = 4 + ratio = 4 + dev = 'cuda' + + q = torch.randn(sq, idx_nh, idx_hd, dtype=torch.float32, device=dev) + k = torch.randn(n_compressed, idx_hd, dtype=torch.float32, device=dev) + w = torch.randn(sq, idx_nh, dtype=torch.float32, device=dev) + cu_q = torch.tensor([0, sq], dtype=torch.int32, device=dev) + cu_kv = torch.tensor([0, n_compressed], dtype=torch.int32, device=dev) + + _, topk_idxs = fused_qk_topk_naive_thd(q, k, w, topk, cu_q, cu_kv, ratio) + + # Causal-valid count per row: min((pos+1)//ratio, n_compressed, topk). + for row in range(sq): + n_valid = min((row + 1) // ratio, n_compressed, topk) + row_idxs = topk_idxs[row] + assert ( + row_idxs[:n_valid] >= 0 + ).all() or n_valid == 0, f"row {row}: leading {n_valid} should be valid ids" + assert (row_idxs[n_valid:] == -1).all(), f"row {row}: tail beyond {n_valid} must be -1" diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py 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_csa_kernels.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_kernels.py new file mode 100644 index 00000000000..a91cb076b26 --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_kernels.py @@ -0,0 +1,3514 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for ``megatron.core.transformer.experimental_attention_variant.csa_kernels``. + +Coverage: + +* Pure-Python helpers: :func:`local_to_global_flat`, :func:`build_flat_topk_idxs`, + :func:`_kl_loss_from_target_predict` — full correctness checks; no GPU + kernels required (CPU is fine). +* Lazy-import gates: :func:`_ensure_flash_mla`, :func:`_ensure_dsa_namespace` + raise informative ``ImportError`` when the optional packages are missing. +* GPU helpers: :func:`_get_topk_alignment` — runs only on CUDA. +* Wrapper functions :func:`_dsa_fwd_flash_mla`, :func:`indexer_topk`, + :func:`dsa_sparse_attn`, :func:`fused_indexer_sparse_attn` — exercised with + ``unittest.mock`` stand-ins for the underlying ``flash_mla`` / + ``cudnn.DSA`` kernels so the data-marshalling logic (shape conversions, + TopK padding, predict/target/KL composition, autograd plumbing) is + validated without requiring the real CUDA kernels. +""" + +from __future__ import annotations + +import math +import sys +import types +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from megatron.core.transformer.experimental_attention_variant import csa_kernels as dk +from megatron.core.transformer.experimental_attention_variant.csa_kernels import ( + FusedIndexerSparseAttnFromTopkFunc, + FusedIndexerSparseAttnFunc, + SparseAttnFunc, + _dsa_fwd_flash_mla, + _ensure_dsa_namespace, + _ensure_flash_mla, + _get_topk_alignment, + _kl_loss_from_dense_scores, + _kl_loss_from_target_predict, + batch_of_row, + build_flat_topk_idxs, + dsa_sparse_attn, + fused_indexer_sparse_attn, + indexer_topk, + local_to_global_flat, +) + +# --------------------------------------------------------------------------- +# Test fixtures / helpers +# --------------------------------------------------------------------------- + + +@pytest.fixture +def reset_lazy_kernel_state(): + """Reset the module-level lazy import slots before/after each test. + + The wrapper-function tests patch ``_flash_mla_sparse_fwd`` / ``_DSA`` + directly; we need to ensure each test starts from a clean slate so the + lazy ``_ensure_*`` calls are exercised consistently. + """ + saved_flash = dk._flash_mla_sparse_fwd + saved_dsa = dk._DSA + dk._flash_mla_sparse_fwd = None + dk._DSA = None + yield + dk._flash_mla_sparse_fwd = saved_flash + dk._DSA = saved_dsa + + +def _make_local_idxs(b: int, sq: int, topk: int, *, with_invalid: bool = False) -> torch.Tensor: + """Build a deterministic ``(b, sq, topk)`` int64 tensor of local indices. + + Values for batch ``i``, query ``s``, slot ``k`` are + ``i * 100 + s * 10 + k``. When ``with_invalid`` is True every other + slot is replaced with -1. + """ + base = ( + torch.arange(b, dtype=torch.int64).view(b, 1, 1) * 100 + + torch.arange(sq, dtype=torch.int64).view(1, sq, 1) * 10 + + torch.arange(topk, dtype=torch.int64).view(1, 1, topk) + ) + if with_invalid: + mask = torch.arange(topk).view(1, 1, topk) % 2 == 1 + base = torch.where(mask.expand(b, sq, topk), torch.full_like(base, -1), base) + return base + + +def _uniform_dist(B, S, K, dev): + """Uniform ``1/K`` distribution of shape ``(B, S, K)``.""" + return torch.full((B, S, K), 1.0 / max(K, 1), dtype=torch.float32, device=dev) + + +def _peaked_dist(B, S, K, dev, peak_idx=0): + """Distribution with all probability mass on ``peak_idx``.""" + out = torch.zeros(B, S, K, dtype=torch.float32, device=dev) + out[..., peak_idx] = 1.0 + return out + + +# --------------------------------------------------------------------------- +# local_to_global_flat +# --------------------------------------------------------------------------- + + +class TestLocalToGlobalFlat: + """Pure-Python index conversion (no GPU required).""" + + @pytest.mark.parametrize( + "b, sq, topk, with_invalid", + [ + (1, 4, 5, False), # b=1 identity case + (2, 3, 4, False), # basic multi-batch + (3, 5, 4, False), # larger batch (stresses the formula) + (2, 3, 6, True), # invalid entries interleaved with valid ones + ], + ids=['b1_identity', 'basic', 'larger_b', 'with_invalid'], + ) + def test_global_index_conversion(self, b, sq, topk, with_invalid): + """Shape, dtype, ``-1`` preservation, and the formula + ``global[s*b + bid, k] = local[bid, s, k] * b + bid`` (for valid entries) + in one fixture. Row ``r`` of the output corresponds to query ``s = r // b`` + and batch id ``bid = r % b``. + """ + local = _make_local_idxs(b, sq, topk, with_invalid=with_invalid) + out = local_to_global_flat(local, b) + + assert out.shape == (sq * b, topk) + assert out.dtype == torch.int32 + + permuted = local.permute(1, 0, 2).reshape(sq * b, topk) + batch_ids = (torch.arange(sq * b) % b).unsqueeze(1) + expected = torch.where( + permuted >= 0, permuted * b + batch_ids, torch.full_like(permuted, -1) + ).int() + assert torch.equal(out, expected) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_cpu_cuda_parity(self): + """CPU and CUDA execution paths produce identical results.""" + local = _make_local_idxs(b=2, sq=4, topk=3, with_invalid=True) + out_cpu = local_to_global_flat(local, 2) + out_cuda = local_to_global_flat(local.cuda(), 2) + assert torch.equal(out_cpu, out_cuda.cpu()) + + +# --------------------------------------------------------------------------- +# build_flat_topk_idxs +# --------------------------------------------------------------------------- + + +class TestBuildFlatTopkIdxs: + """Pure-Python multi-group concat + optional compaction.""" + + @pytest.mark.parametrize( + "group_specs", + [ + # Each spec is a list of (topk_i, with_invalid) for each group. + [(4, False)], # single group, all valid + [(2, False), (3, False)], # two groups, all valid + ], + ids=['single_group', 'two_groups'], + ) + def test_non_compact_concat_then_globalise(self, group_specs): + """Without ``compact`` the helper concatenates groups along ``topk`` + and applies the local→global conversion verbatim. + """ + b, sq = 2, 3 + groups = [ + _make_local_idxs(b, sq, t, with_invalid=inv) + 50 * i + for i, (t, inv) in enumerate(group_specs) + ] + total_topk = sum(t for t, _ in group_specs) + + flat, length = build_flat_topk_idxs(*groups, batch_size=b) + + expected = local_to_global_flat(torch.cat(groups, dim=-1), b) + assert flat.shape == (sq * b, total_topk) + assert flat.dtype == torch.int32 + assert torch.equal(flat, expected) + assert length is None + + @pytest.mark.parametrize( + "group_specs, expected_valid_per_row", + [ + # Single group: 6 slots with every odd slot invalid → 3 valid. + ([(6, True)], 3), + # Two groups: g1 has 2 valid out of 4, g2 fully valid (2) → 4 valid. + ([(4, True), (2, False)], 4), + ], + ids=['single_group', 'two_groups'], + ) + def test_compact_packs_valid_first(self, group_specs, expected_valid_per_row): + """With ``compact=True`` the helper packs valid entries to the front + of each row, fills the tail with ``-1``, and returns a per-row + ``topk_length`` that equals the count of valid entries. + """ + b, sq = 2, 3 + groups = [ + _make_local_idxs(b, sq, t, with_invalid=inv) + 100 * i + for i, (t, inv) in enumerate(group_specs) + ] + total_topk = sum(t for t, _ in group_specs) + + flat, length = build_flat_topk_idxs(*groups, batch_size=b, compact=True) + + assert flat.shape == (sq * b, total_topk) + assert flat.dtype == torch.int32 + assert length is not None + assert length.shape == (sq * b,) + assert length.dtype == torch.int32 + + # Per-row layout: valid global indices first, then -1 padding. + for row in range(sq * b): + n = int(length[row]) + assert n == expected_valid_per_row, f"row {row}: wrong length" + assert torch.all(flat[row, :n] >= 0), f"row {row}: leading entries should be valid" + assert torch.all(flat[row, n:] == -1), f"row {row}: trailing entries should be -1" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_compact_cuda_path(self, reset_lazy_kernel_state): + """Combined coverage for the CUDA compact path (sub-blocks + self-label on failure): + + * (a) Dispatch + plumbing (mocked compactify): the wrapper is + called exactly once, with the already-globalised + ``(sq*b, total_topk)`` int32 tensor as input, and its + returned ``(indices, topk_length)`` flow back verbatim. + * (b) End-to-end parity (real cuDNN, skipped without it): the + cuDNN ``compactify`` kernel produces the same ``(flat, + length)`` pair as the pure-PyTorch CPU fallback. + """ + # ---- (a) dispatch via mocked compactify -------------------------- + b, sq, topk = 2, 3, 4 + local = _make_local_idxs(b, sq, topk, with_invalid=True).to(torch.int32, copy=False).cuda() + compact_indices = torch.full((sq * b, topk), 99, dtype=torch.int32, device='cuda') + topk_length = torch.full((sq * b,), 7, dtype=torch.int32, device='cuda') + + captured = {} + + def fake_compactify(global_idxs): + captured['input'] = global_idxs + return {'indices': compact_indices, 'topk_length': topk_length} + + fake_dsa = MagicMock(name='_DSA_compactify_stub') + fake_dsa.compactify_wrapper.side_effect = fake_compactify + dk._DSA = fake_dsa + + flat, length = build_flat_topk_idxs(local, batch_size=b, compact=True) + fake_dsa.compactify_wrapper.assert_called_once() + kernel_input = captured['input'] + assert kernel_input.shape == (sq * b, topk), "(a) wrapper input shape" + assert kernel_input.dtype == torch.int32, "(a) wrapper input dtype" + assert kernel_input.is_cuda, "(a) wrapper input not on CUDA" + expected_input = local_to_global_flat(local, b) + assert torch.equal( + kernel_input, expected_input + ), "(a) wrapper input != local_to_global_flat(local)" + assert flat is compact_indices, "(a) returned flat is not the kernel output" + assert length is topk_length, "(a) returned length is not the kernel output" + + # ---- (b) real-kernel parity vs CPU fallback ---------------------- + # Skipped when cuDNN is not installed; reset state so the real + # _DSA import happens on the next call inside build_flat_topk_idxs. + try: + cudnn = pytest.importorskip("cudnn") + except pytest.skip.Exception: + return # already passed (a); skip the parity sub-block silently + if not hasattr(cudnn, 'DSA'): + return + dk._DSA = None # force real lazy-import + + b2, sq2 = 4, 5 + local_a = _make_local_idxs(b2, sq2, 6, with_invalid=True) + local_b = _make_local_idxs(b2, sq2, 4, with_invalid=False) + 200 + + flat_cpu, len_cpu = build_flat_topk_idxs(local_a, local_b, batch_size=b2, compact=True) + flat_cuda, len_cuda = build_flat_topk_idxs( + local_a.cuda(), local_b.cuda(), batch_size=b2, compact=True + ) + assert torch.equal( + flat_cpu, flat_cuda.cpu() + ), "(b) flat tensor differs between CPU fallback and cuDNN kernel" + assert torch.equal( + len_cpu, len_cuda.cpu() + ), "(b) length tensor differs between CPU fallback and cuDNN kernel" + + +# --------------------------------------------------------------------------- +# _kl_loss_from_target_predict +# --------------------------------------------------------------------------- + + +class TestKLLossFromTargetPredict: + """Pure-Python KL loss computation: combined assertions for all + properties (scalar/dtype, identity, non-negativity, coeff linearity, + invalid-row masking, analytical formula).""" + + def test_kl_loss_properties(self): + """All KL-loss invariants checked sequentially. Each block raises + an informative ``AssertionError`` so a failure pinpoints the + broken sub-property. + """ + torch.manual_seed(0) + b, sq, topk = 2, 3, 4 + topk_indices = torch.zeros(b, sq, topk, dtype=torch.int32) + + # ---- (a) scalar/dtype + identity: KL(p || p) == 0 ----------------- + identical = torch.softmax(torch.randn(b, sq, topk), dim=-1) + loss_identical = _kl_loss_from_target_predict( + identical, identical.clone(), topk_indices, loss_coeff=1.0 + ) + assert loss_identical.shape == torch.Size([]), "identity: not scalar" + assert loss_identical.dtype == torch.float32, "identity: not fp32" + assert torch.allclose( + loss_identical, torch.tensor(0.0), atol=1e-6 + ), f"identity: KL(p || p) != 0 (got {loss_identical.item()})" + + # ---- (b) non-negativity + linearity in loss_coeff ----------------- + target = torch.softmax(torch.randn(b, sq, topk), dim=-1) + predict = torch.softmax(torch.randn(b, sq, topk), dim=-1) + loss_1 = _kl_loss_from_target_predict(target, predict, topk_indices, loss_coeff=1.0) + loss_3 = _kl_loss_from_target_predict(target, predict, topk_indices, loss_coeff=3.0) + assert loss_1.item() >= 0.0, f"non-negativity: got {loss_1.item()}" + assert torch.allclose( + loss_3, 3.0 * loss_1, atol=1e-5, rtol=1e-5 + ), f"linearity: 3*loss_1 = {3*loss_1.item()} vs loss_3 = {loss_3.item()}" + + # ---- (c) invalid-row masking -------------------------------------- + # Construct deterministic distributions with strictly-positive per-row KL. + t_inv = torch.full((b, sq, topk), 0.1, dtype=torch.float32) + t_inv[..., 0] = 0.7 + p_inv = torch.full((b, sq, topk), 0.7, dtype=torch.float32) / topk + p_inv[..., -1] = 1.0 - p_inv[..., :-1].sum(dim=-1) + + idx_all_valid = torch.zeros(b, sq, topk, dtype=torch.int32) + loss_full = _kl_loss_from_target_predict(t_inv, p_inv, idx_all_valid, loss_coeff=1.0) + assert loss_full.item() > 0, "all-valid baseline must be positive" + + # Mark the first row of every batch invalid → fewer valid rows, + # smaller KL sum, same denominator (mean over all (B, S_q)). + idx_partial = idx_all_valid.clone() + idx_partial[:, 0, :] = -1 + loss_partial = _kl_loss_from_target_predict(t_inv, p_inv, idx_partial, loss_coeff=1.0) + assert ( + loss_partial.item() < loss_full.item() + ), f"partial-invalid: {loss_partial.item()} should be < {loss_full.item()}" + + # All-invalid → loss exactly 0. + idx_all_invalid = torch.full_like(idx_all_valid, -1) + loss_zero = _kl_loss_from_target_predict(t_inv, p_inv, idx_all_invalid, loss_coeff=1.0) + assert loss_zero.item() == 0.0, f"all-invalid: got {loss_zero.item()}" + + # ---- (d) analytical formula: target = δ_0, predict = uniform(1/K) - + # per-row KL = log(K); mean = log(K); loss = coeff * log(K). + target_d = _peaked_dist(b, sq, topk, 'cpu', peak_idx=0) + predict_d = torch.full((b, sq, topk), 1.0 / topk, dtype=torch.float32) + loss_d = _kl_loss_from_target_predict(target_d, predict_d, topk_indices, loss_coeff=2.5) + expected = 2.5 * math.log(topk) + assert torch.allclose( + loss_d, torch.tensor(expected), rtol=1e-5, atol=1e-5 + ), f"analytical: {loss_d.item()} vs expected {expected}" + + def test_per_token_loss_reports_raw_sum(self): + torch.manual_seed(1) + b, sq, topk = 2, 5, 4 + target = torch.softmax(torch.randn(b, sq, topk), dim=-1) + predict = torch.softmax(torch.randn(b, sq, topk), dim=-1) + topk_indices = torch.zeros(b, sq, topk, dtype=torch.int32) + + loss_mean = _kl_loss_from_target_predict(target, predict, topk_indices, loss_coeff=0.5) + loss_sum = _kl_loss_from_target_predict( + target, predict, topk_indices, loss_coeff=0.5, calculate_per_token_loss=True + ) + + assert torch.allclose(loss_sum, loss_mean * (b * sq), rtol=1e-5, atol=1e-5) + + +class TestKLLossFromDenseScores: + def test_per_token_loss_reports_raw_sum(self): + b, sq, sk = 2, 5, 4 + loss_coeff = 0.5 + + attn_score = _peaked_dist(b, sq, sk, 'cpu', peak_idx=0) + attn_l1norm = torch.ones(b, sq, dtype=torch.float32) + index_score = torch.zeros(b, sq, sk, dtype=torch.float32) + index_lse = torch.full((b, sq), math.log(sk), dtype=torch.float32) + + loss_mean = _kl_loss_from_dense_scores( + attn_score, attn_l1norm, index_score, index_lse, loss_coeff + ) + loss_sum = _kl_loss_from_dense_scores( + attn_score, + attn_l1norm, + index_score, + index_lse, + loss_coeff, + calculate_per_token_loss=True, + ) + + assert torch.allclose(loss_sum, loss_mean * (b * sq), rtol=1e-5, atol=1e-5) + + def test_ratio_masked_positions_do_not_produce_infinite_loss(self): + attn_score = torch.tensor([[1.0, 0.0, 0.0]]) + attn_l1norm = torch.tensor([1.0]) + index_score = torch.tensor([[0.0, float("-inf"), float("-inf")]]) + index_lse = torch.tensor([0.0]) + + loss = _kl_loss_from_dense_scores( + attn_score, attn_l1norm, index_score, index_lse, loss_coeff=1.0 + ) + + torch.testing.assert_close(loss, torch.zeros_like(loss)) + + +# --------------------------------------------------------------------------- +# _ensure_flash_mla / _ensure_dsa_namespace +# --------------------------------------------------------------------------- + + +_LAZY_IMPORT_CASES = [ + pytest.param( + 'flash_mla', + 'flash_mla_sparse_fwd', + _ensure_flash_mla, + '_flash_mla_sparse_fwd', + "FlashMLA is required", + id='flash_mla', + ), + pytest.param( + 'cudnn', 'DSA', _ensure_dsa_namespace, '_DSA', "cudnn-frontend DSA", id='cudnn_dsa' + ), +] + + +@pytest.mark.parametrize( + "module_name, attr_name, ensure_fn, slot_name, error_match", _LAZY_IMPORT_CASES +) +class TestLazyKernelImports: + """Lazy-import behaviour shared by ``_ensure_flash_mla`` and + ``_ensure_dsa_namespace``: combined error-on-missing + caches-on-success + fixture (assertion blocks self-label on failure). + """ + + def test_lazy_import_raises_and_caches( + self, reset_lazy_kernel_state, module_name, attr_name, ensure_fn, slot_name, error_match + ): + # ---- (a) raises informative ImportError when the module is absent -- + # Setting ``sys.modules[name] = None`` makes ``import name`` fail. + with patch.dict(sys.modules, {module_name: None}): + with pytest.raises(ImportError, match=error_match): + ensure_fn() + + # ---- (b) caches the import on success ------------------------------ + sentinel = MagicMock(name=f"{attr_name}_sentinel") + fake_module = types.ModuleType(module_name) + setattr(fake_module, attr_name, sentinel) + + with patch.dict(sys.modules, {module_name: fake_module}): + ensure_fn() + assert ( + getattr(dk, slot_name) is sentinel + ), f"(b) {module_name}: ensure_fn() did not bind sentinel" + + # Second call must be a no-op — even after the sys.modules entry is gone. + with patch.dict(sys.modules, {}, clear=False): + sys.modules.pop(module_name, None) + ensure_fn() + assert ( + getattr(dk, slot_name) is sentinel + ), f"(b) {module_name}: cached sentinel was lost on 2nd call" + + +# --------------------------------------------------------------------------- +# _get_topk_alignment +# --------------------------------------------------------------------------- + + +class TestGetTopkAlignment: + """Architecture-dependent alignment for FlashMLA top-K padding.""" + + @pytest.fixture(autouse=True) + def _clear_alignment_cache(self): + # ``_get_topk_alignment`` is ``@lru_cache``-d, so the first call freezes + # its result for the process. Clear it around every test so the patched + # device capability is actually re-read. + _get_topk_alignment.cache_clear() + yield + _get_topk_alignment.cache_clear() + + @pytest.mark.parametrize( + "sm_major, expected", [(7, 128), (8, 128), (9, 128), (10, 64), (12, 64), (13, 64)] + ) + def test_alignment_per_sm(self, sm_major, expected): + """SM10x and newer use 64-byte TopK alignment; older arches use 128.""" + with patch('torch.cuda.get_device_capability', return_value=(sm_major, 0)): + assert _get_topk_alignment() == expected + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_runs_on_real_gpu(self): + """On any real GPU the alignment must agree with the documented rule.""" + align = _get_topk_alignment() + sm = torch.cuda.get_device_capability() + expected = 64 if sm[0] >= 10 else 128 + assert align == expected + + +# --------------------------------------------------------------------------- +# _dsa_fwd_flash_mla — wrapper around flash_mla.flash_mla_sparse_fwd +# --------------------------------------------------------------------------- + + +def _make_flash_mla_stub(d_v: int = 512, *, lse_scalar: float = 0.0, out_fill: float = 0.0): + """Build a callable stand-in for ``flash_mla.flash_mla_sparse_fwd``. + + The real kernel signature is + ``(q, kv, indices, softmax_scale, d_v, attn_sink, topk_length, indexer_topk)`` + and returns ``(out, max_logits, lse)`` or ``(out, max_logits, lse, lse_indexer)`` + when ``indexer_topk > 0``. + + The stub returns deterministic, easily-distinguishable tensors so callers + can numerically verify the wrapper's reshape / split logic. The most + recent ``out`` and ``lse`` are stashed on ``stub.last_out`` / + ``stub.last_lse`` for direct equality checks. + """ + + stub = MagicMock(name='flash_mla_sparse_fwd_stub') + + def _impl(q, kv, indices, softmax_scale, d_v, attn_sink, topk_length, indexer_topk): + total_S_q, H, _D = q.shape + # Distinguishable per-element pattern: out[i, h, k] = out_fill + i + 0.001*h + 1e-6*k + # (works in bf16 at this magnitude, useful for verifying that the + # wrapper does not silently reshape across the wrong axes). + idx_i = torch.arange(total_S_q, dtype=torch.float32, device=q.device).view(-1, 1, 1) + idx_h = torch.arange(H, dtype=torch.float32, device=q.device).view(1, -1, 1) + idx_k = torch.arange(d_v, dtype=torch.float32, device=q.device).view(1, 1, -1) + out_f32 = out_fill + idx_i + 0.001 * idx_h + 1e-6 * idx_k + out = out_f32.to(q.dtype) + max_logits = torch.zeros(total_S_q, H, dtype=torch.float32, device=q.device) + # lse[i, h] = lse_scalar + i + 0.5*h — a deterministic pattern. + lse = lse_scalar + ( + torch.arange(total_S_q, dtype=torch.float32, device=q.device).view(-1, 1) + + 0.5 * torch.arange(H, dtype=torch.float32, device=q.device).view(1, -1) + ) + stub.last_out = out + stub.last_lse = lse + if indexer_topk > 0: + # Make lse_indexer distinct from lse so we can tell which one the + # wrapper returned. + lse_indexer = lse + 100.0 + stub.last_lse_indexer = lse_indexer + return out, max_logits, lse, lse_indexer + stub.last_lse_indexer = None + return out, max_logits, lse + + stub.side_effect = _impl + stub.last_out = None + stub.last_lse = None + stub.last_lse_indexer = None + return stub + + +class TestDsaFwdFlashMla: + """Adapter logic around FlashMLA: shape massaging, TopK padding, return + tuples — including numerical pass-through of the kernel outputs. + """ + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dsa_fwd_flash_mla_adapter(self, reset_lazy_kernel_state): + """All adapter behaviours in one fixture (assertion blocks self-label + on failure): + + * (a) TopK is padded up to GPU-specific alignment; padded slots are + ``-1``; ``out`` / ``lse`` are passed through verbatim; + kernel arg shapes match the SBHD-flat → MQA-h_kv=1 contract. + * (b) ``indexer_topk == 0`` -> ``lse_indexer is None``. + * (c) ``0 < indexer_topk < TopK`` -> kernel's ``lse_indexer`` is + returned verbatim (not silently swapped for ``lse``). + * (d) ``indexer_topk == TopK`` -> fallback to ``lse.clone()`` (a + kernel snapshot quirk). + * (e) ``indexer_topk > 0 + topk_length`` is rejected with the + expected error. + """ + total_sq, H, D = 4, 2, 512 + align = _get_topk_alignment() + + q = torch.randn(total_sq, H, D, dtype=torch.bfloat16, device='cuda') + kv = torch.randn(8, D, dtype=torch.bfloat16, device='cuda') + + # ---- (a) padding + numerical pass-through ------------------------ + TopK_unpadded = 5 + expected_padded = ((TopK_unpadded + align - 1) // align) * align + topk_idxs = torch.arange(total_sq * TopK_unpadded, dtype=torch.int32, device='cuda').view( + total_sq, TopK_unpadded + ) + + stub = _make_flash_mla_stub(d_v=D) + dk._flash_mla_sparse_fwd = stub + + out, lse, lse_indexer = _dsa_fwd_flash_mla(q, kv, topk_idxs, softmax_scale=0.5, d_v=D) + assert lse_indexer is None, "(a) lse_indexer should be None when indexer_topk=0" + assert torch.equal(out, stub.last_out), "(a) out is not pass-through" + assert torch.equal(lse, stub.last_lse), "(a) lse is not pass-through" + + called_args = stub.call_args.args + kv_3d, indices_arg = called_args[1], called_args[2] + assert kv_3d.shape == (8, 1, D), f"(a) KV shape {tuple(kv_3d.shape)} != (8, 1, {D})" + assert indices_arg.shape == (total_sq, 1, expected_padded), ( + f"(a) indices shape {tuple(indices_arg.shape)} != " + f"({total_sq}, 1, {expected_padded})" + ) + if expected_padded > TopK_unpadded: + assert torch.all( + indices_arg[..., TopK_unpadded:] == -1 + ), "(a) padded slots should be -1" + assert torch.equal( + indices_arg[..., :TopK_unpadded].squeeze(1), topk_idxs + ), "(a) original entries should survive padding unchanged" + + # ---- (b–d) indexer_topk branches --------------------------------- + TopK = align # already aligned, no padding + topk_idxs_aligned = torch.zeros(total_sq, TopK, dtype=torch.int32, device='cuda') + + stub = _make_flash_mla_stub(d_v=D, lse_scalar=1.5) + dk._flash_mla_sparse_fwd = stub + + # (b) indexer_topk == 0 + _, _, lse_idx_b = _dsa_fwd_flash_mla(q, kv, topk_idxs_aligned, 0.5, indexer_topk=0) + assert lse_idx_b is None, "(b) indexer_topk=0 must yield lse_indexer=None" + + # (c) 0 < indexer_topk < TopK + _, lse_c, lse_idx_c = _dsa_fwd_flash_mla( + q, kv, topk_idxs_aligned, 0.5, indexer_topk=TopK // 2 + ) + assert lse_idx_c is not None, "(c) lse_indexer should be present" + assert torch.equal( + lse_idx_c, stub.last_lse_indexer + ), "(c) lse_indexer should be kernel pass-through" + assert not torch.equal(lse_idx_c, lse_c), "(c) wrapper silently swapped lse_indexer for lse" + + # (d) indexer_topk == TopK -> fallback to lse.clone() + _, lse_d, lse_idx_d = _dsa_fwd_flash_mla(q, kv, topk_idxs_aligned, 0.5, indexer_topk=TopK) + assert torch.equal( + lse_idx_d, lse_d + ), "(d) lse_indexer should equal lse on TopK-cap fallback" + assert ( + lse_idx_d.data_ptr() != lse_d.data_ptr() + ), "(d) fallback should be a clone, not an alias" + + # ---- (e) topk_length + indexer_topk > 0 is rejected -------------- + # Use CPU tensors here — the assert fires before any kernel call. + with pytest.raises(AssertionError, match="indexer_topk > 0 requires non-compact"): + _dsa_fwd_flash_mla( + torch.zeros(2, 2, 512, dtype=torch.bfloat16), + torch.zeros(4, 512, dtype=torch.bfloat16), + torch.zeros(2, 8, dtype=torch.int32), + softmax_scale=0.5, + topk_length=torch.zeros(2, dtype=torch.int32), + indexer_topk=4, + ) + + +# --------------------------------------------------------------------------- +# indexer_topk — cudnn DSA wrapper for inference +# --------------------------------------------------------------------------- + + +class TestIndexerTopk: + """Indexer scoring + radix top-K wrapper. All three properties combined + in a single fixture; sub-block names appear in failure messages. + """ + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_indexer_topk_wrapper(self, reset_lazy_kernel_state): + """Combined assertions for: + + * (a) basic call: shapes / dtypes of the returned (topk_indices, + topk_length); kernels are called with the right BSHD layouts + and the SBHD-flat (b*sq, sk) scores; indexer_top_k kwargs. + * (b) topk > sk clamping: kernel call uses ``sk`` keys, trailing + ``[sk:]`` slots are -1, ``topk_length == sk``. + * (c) ``indexer_softmax_scale`` pre-scales the weights via the + ``relu(c·x) = c·relu(x)`` trick before reaching the kernel. + """ + + # ---- (a) basic call ---------------------------------------------- + sq, b, idx_nh, idx_hd = 6, 2, 4, 64 + sk = 12 + topk = 5 + ratio = 4 + + q_indexer = torch.randn(sq, b, idx_nh, idx_hd, dtype=torch.bfloat16, device='cuda') + k_indexer = torch.randn(sk, b, idx_hd, dtype=torch.bfloat16, device='cuda') + weights = torch.randn(sq, b, idx_nh, dtype=torch.bfloat16, device='cuda') + + scores = torch.randn(b, sq, sk, dtype=torch.float32, device='cuda') + captured = {} + + def fake_indexer_forward(q_bshd, k_bshd, w_bsh, ratio): + captured['indexer_forward'] = { + 'q_shape': q_bshd.shape, + 'k_shape': k_bshd.shape, + 'w_shape': w_bsh.shape, + 'ratio': ratio, + } + return {'scores': scores} + + def fake_filtered_topk(scores_flat, seq_lens, top_k, next_n, return_val): + captured['filtered_topk'] = { + 'scores_shape': scores_flat.shape, + 'seq_lens_shape': seq_lens.shape, + 'top_k': top_k, + 'next_n': next_n, + 'return_val': return_val, + } + n_rows = scores_flat.shape[0] + return {'indices': torch.zeros(n_rows, top_k, dtype=torch.int32, device='cuda')} + + fake_dsa = MagicMock() + fake_dsa.indexer_forward_wrapper.side_effect = fake_indexer_forward + fake_dsa.indexer_top_k_wrapper.side_effect = fake_filtered_topk + dk._DSA = fake_dsa + + topk_indices, topk_length = indexer_topk( + q_indexer, k_indexer, weights, topk=topk, ratio=ratio + ) + + assert topk_indices.shape == (b, sq, topk), "(a) topk_indices shape" + assert topk_indices.dtype == torch.int32, "(a) topk_indices dtype" + assert topk_length.shape == (b, sq), "(a) topk_length shape" + assert topk_length.dtype == torch.int32, "(a) topk_length dtype" + # BSHD / BSD layouts handed to the kernels. + assert captured['indexer_forward']['q_shape'] == ( + b, + sq, + idx_nh, + idx_hd, + ), "(a) indexer_forward q_shape" + assert captured['indexer_forward']['k_shape'] == ( + b, + sk, + 1, + idx_hd, + ), "(a) indexer_forward k_shape (must be unsqueezed h_kv=1)" + assert captured['indexer_forward']['w_shape'] == ( + b, + sq, + idx_nh, + ), "(a) indexer_forward w_shape" + assert captured['indexer_forward']['ratio'] == ratio, "(a) ratio kwarg" + assert captured['filtered_topk']['scores_shape'] == ( + b * sq, + sk, + ), "(a) topK scores_flat shape" + assert captured['filtered_topk']['seq_lens_shape'] == (b * sq,), "(a) topK seq_lens shape" + assert captured['filtered_topk']['top_k'] == min(topk, sk), "(a) top_k kwarg" + assert captured['filtered_topk']['next_n'] == 1, "(a) next_n kwarg" + assert captured['filtered_topk']['return_val'] is False, "(a) return_val kwarg" + + # ---- (b) topk > sk clamping -------------------------------------- + dk._DSA = None # force fresh mocks + sq2, b2, idx_nh2, idx_hd2 = 4, 1, 2, 32 + sk2 = 3 + topk2 = 8 # > sk + q2 = torch.randn(sq2, b2, idx_nh2, idx_hd2, dtype=torch.bfloat16, device='cuda') + k2 = torch.randn(sk2, b2, idx_hd2, dtype=torch.bfloat16, device='cuda') + w2 = torch.randn(sq2, b2, idx_nh2, dtype=torch.bfloat16, device='cuda') + scores2 = torch.zeros(b2, sq2, sk2, dtype=torch.float32, device='cuda') + kernel_indices2 = torch.zeros(b2 * sq2, sk2, dtype=torch.int32, device='cuda') + + fake_dsa_b = MagicMock() + fake_dsa_b.indexer_forward_wrapper.return_value = {'scores': scores2} + fake_dsa_b.indexer_top_k_wrapper.return_value = {'indices': kernel_indices2} + dk._DSA = fake_dsa_b + + topk_indices2, topk_length2 = indexer_topk(q2, k2, w2, topk=topk2, ratio=4) + assert topk_indices2.shape == (b2, sq2, topk2), "(b) padded topk_indices shape" + assert torch.all(topk_indices2[..., sk2:] == -1), "(b) trailing slots not -1" + assert torch.all(topk_length2 == sk2), "(b) topk_length should equal sk" + + # ---- (c) indexer_softmax_scale pre-scales weights --------------- + dk._DSA = None + sq3, b3, idx_nh3, idx_hd3 = 2, 1, 2, 32 + sk3 = 4 + scale = 0.125 + q3 = torch.zeros(sq3, b3, idx_nh3, idx_hd3, dtype=torch.bfloat16, device='cuda') + k3 = torch.zeros(sk3, b3, idx_hd3, dtype=torch.bfloat16, device='cuda') + w3 = torch.full((sq3, b3, idx_nh3), 8.0, dtype=torch.bfloat16, device='cuda') + captured_w = {} + + def fake_indexer_forward_c(q_bshd, k_bshd, w_bsh, ratio): + captured_w['w'] = w_bsh.detach().clone() + return {'scores': torch.zeros(b3, sq3, sk3, dtype=torch.float32, device='cuda')} + + fake_dsa_c = MagicMock() + fake_dsa_c.indexer_forward_wrapper.side_effect = fake_indexer_forward_c + fake_dsa_c.indexer_top_k_wrapper.return_value = { + 'indices': torch.zeros(b3 * sq3, sk3, dtype=torch.int32, device='cuda') + } + dk._DSA = fake_dsa_c + + indexer_topk(q3, k3, w3, topk=sk3, ratio=4, indexer_softmax_scale=scale) + expected_w = ( + (w3.float() * scale).to(torch.bfloat16).permute(1, 0, 2).reshape(b3, sq3, idx_nh3) + ) + assert torch.allclose( + captured_w['w'].float(), expected_w.float(), atol=1e-2, rtol=1e-2 + ), "(c) weights were not pre-scaled by indexer_softmax_scale" + + +# --------------------------------------------------------------------------- +# dsa_sparse_attn / SparseAttnFunc forward (mocked) +# --------------------------------------------------------------------------- + + +class TestDsaSparseAttn: + """Numerical fwd + bwd test for the public ``dsa_sparse_attn`` entry + point. The underlying kernels are mocked so the whole wrapper — including + the SBHD↔flat reshape on the forward and the autograd plumbing on the + backward — can be checked against deterministic ground truth. + """ + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dsa_sparse_attn_fwd_and_bwd(self, reset_lazy_kernel_state): + """Combined forward + backward fixture (assertion blocks self-label + on failure): + + * (a) Forward output equals the FlashMLA stub's ``out`` reshaped + from flat ``(sq*b, np_, d_v)`` back to ``(sq, b, np_ * d_v)``. + * (b) Backward maps kernel grads onto the right SBHD leaf tensors + with the correct shapes; the bwd kernel is invoked exactly + once. + """ + sq, b, np_, d = 4, 2, 2, 512 + skv = 6 + TopK = _get_topk_alignment() + + query = torch.randn(sq, b, np_, d, dtype=torch.bfloat16, device='cuda', requires_grad=True) + kv = torch.randn(skv, b, d, dtype=torch.bfloat16, device='cuda', requires_grad=True) + attn_sink = torch.zeros(np_, dtype=torch.float32, device='cuda', requires_grad=True) + topk_idxs = torch.zeros(sq * b, TopK, dtype=torch.int32, device='cuda') + + # Coordinated stubs: FlashMLA fwd + cuDNN sparse-attn bwd, both + # deterministic so every gradient slot is independently verifiable. + flash_stub = _make_flash_mla_stub(d_v=d) + dq_kernel = torch.full((sq * b, np_, d), 7.0, dtype=torch.bfloat16, device='cuda') + dkv_kernel = torch.full((skv * b, d), -3.0, dtype=torch.bfloat16, device='cuda') + d_sink_kernel = torch.full((np_,), 11.0, dtype=torch.float32, device='cuda') + fake_dsa = MagicMock() + fake_dsa.sparse_attention_backward_wrapper.return_value = { + 'dq': dq_kernel, + 'dkv': dkv_kernel, + 'd_sink': d_sink_kernel, + } + dk._flash_mla_sparse_fwd = flash_stub + dk._DSA = fake_dsa + + out = dsa_sparse_attn(query, kv, attn_sink, topk_idxs, softmax_scale=0.5) + + # ---- (a) forward ------------------------------------------------ + assert out.shape == (sq, b, np_ * d), "(a) forward shape" + assert out.dtype == torch.bfloat16, "(a) forward dtype" + expected_out = flash_stub.last_out.reshape(sq, b, np_, d).reshape(sq, b, np_ * d) + assert torch.equal(out, expected_out), "(a) forward value pass-through" + + # ---- (b) backward ----------------------------------------------- + out.sum().backward() + assert query.grad is not None, "(b) query.grad missing" + assert kv.grad is not None, "(b) kv.grad missing" + assert attn_sink.grad is not None, "(b) attn_sink.grad missing" + assert torch.equal( + query.grad, dq_kernel.reshape(sq, b, np_, d) + ), "(b) query.grad mis-reshaped" + assert torch.equal(kv.grad, dkv_kernel.reshape(skv, b, d)), "(b) kv.grad mis-reshaped" + assert torch.equal(attn_sink.grad, d_sink_kernel), "(b) attn_sink.grad mismatch" + fake_dsa.sparse_attention_backward_wrapper.assert_called_once() + + +# --------------------------------------------------------------------------- +# fused_indexer_sparse_attn — Path B autograd Function (mocked) +# --------------------------------------------------------------------------- + + +def _install_full_dsa_mock( + *, + b: int, + sq: int, + np_: int, + d: int, + n_comp: int, + idx_nh: int, + predict_fn=None, + target_fn=None, + dq_value: float = 7.0, + dkv_value: float = -3.0, + d_sink_value: float = 11.0, + d_index_q_value: float = 0.5, + d_weights_value: float = -0.25, + d_index_k_value: float = 1.5, +): + """Patch the module-level ``_DSA`` and ``_flash_mla_sparse_fwd`` slots + with a coordinated set of deterministic stubs covering every kernel + invoked by :class:`FusedIndexerSparseAttnFunc`. + + ``predict_fn`` / ``target_fn`` (if provided) build the per-row + distribution given ``(b, sq, topk, device)``. By default both return a + uniform ``1/topk`` distribution, which yields ``KL(target || predict) = 0`` + so the loss is exactly zero. + + All backward kernels return constant-filled tensors so each gradient slot + can be independently verified. + """ + + if predict_fn is None: + predict_fn = lambda B, S, K, dev: torch.full( + (B, S, K), 1.0 / max(K, 1), dtype=torch.float32, device=dev + ) + if target_fn is None: + target_fn = predict_fn + + fake_dsa = MagicMock(name='_DSA_full_stub') + + def fake_indexer_forward(q_bshd, k_bshd, w_bsh, ratio): + return {'scores': torch.zeros(b, sq, n_comp, dtype=torch.float32, device=q_bshd.device)} + + fake_dsa.indexer_forward_wrapper.side_effect = fake_indexer_forward + + def fake_filtered_topk(scores_flat, seq_lens, top_k, next_n, return_val): + return { + 'indices': torch.zeros( + scores_flat.shape[0], top_k, dtype=torch.int32, device=scores_flat.device + ) + } + + fake_dsa.indexer_top_k_wrapper.side_effect = fake_filtered_topk + + def fake_sparse_indexer_score_backward( + q, k, w, topk_indices, qhead_per_kv_head, topk_indices_global=False + ): + topk = topk_indices.shape[-1] + return {'predict': predict_fn(b, sq, topk, q.device)} + + fake_dsa.sparse_indexer_score_recompute_wrapper.side_effect = fake_sparse_indexer_score_backward + + def fake_sparse_attn_score_backward( + q, k, lse, topk_indices, sm_scale, qhead_per_kv_head, topk_indices_global=False + ): + topk = topk_indices.shape[-1] + return {'target': target_fn(b, sq, topk, q.device)} + + fake_dsa.sparse_attn_score_recompute_wrapper.side_effect = fake_sparse_attn_score_backward + + def fake_sparse_attn_backward(q, kv, out, dout, lse, attn_sink, topk_idxs, **kwargs): + return { + 'dq': torch.full_like(q, dq_value), + 'dkv': torch.full_like(kv, dkv_value), + 'd_sink': torch.full_like(attn_sink, d_sink_value), + } + + fake_dsa.sparse_attention_backward_wrapper.side_effect = fake_sparse_attn_backward + + def fake_indexer_grad_backward( + q_idx_bshd, + w_bsh, + k_idx_bsd, + attn_score, + index_score, + topk_indices, + sm_scale, + loss_coeff, + grad_loss, + block_I, + ): + return { + 'd_index_q': torch.full_like(q_idx_bshd, d_index_q_value), + 'd_weights': torch.full_like(w_bsh, d_weights_value), + 'd_index_k': torch.full_like(k_idx_bsd, d_index_k_value), + } + + fake_dsa.indexer_backward_wrapper.side_effect = fake_indexer_grad_backward + + flash_stub = _make_flash_mla_stub(d_v=d) + + dk._DSA = fake_dsa + dk._flash_mla_sparse_fwd = flash_stub + return fake_dsa, flash_stub + + +def _install_full_dsa_mock_dense( + *, + b: int, + sq: int, + np_: int, + d: int, + n_comp: int, + idx_nh: int, + target_score_fn=None, + target_l1norm_fn=None, + predict_score_fn=None, + predict_lse_fn=None, + dq_value: float = 7.0, + dkv_value: float = -3.0, + d_sink_value: float = 11.0, + d_index_q_value: float = 0.5, + d_weights_value: float = -0.25, + d_index_k_value: float = 1.5, +): + """Coordinated stubs covering the dense-loss (``sparse_loss=False``) path. + + Mirrors :func:`_install_full_dsa_mock` for the sparse path, but stubs + the four dense-only kernel wrappers: + + * ``dense_indexer_score_recompute_wrapper`` -> ``(out, denom=index_lse)`` + * ``dense_attn_score_recompute_wrapper`` -> ``(out, denom=attn_l1norm)`` + * ``dense_indexer_backward_wrapper`` -> ``{d_index_q, d_weights, d_index_k}`` + + Defaults make ``target == predict == uniform(1/n_comp)`` so KL == 0. + Override the four ``*_fn`` callables to drive the loss to known + analytical values; each callable receives ``(B, S_q, S_k, device)`` and + returns the score tensor (``S_k``-dim) or denom (no ``S_k`` dim). + """ + + if target_score_fn is None: + target_score_fn = lambda B, S, K, dev: torch.full( + (B, S, K), 1.0 / max(K, 1), dtype=torch.float32, device=dev + ) + if target_l1norm_fn is None: + target_l1norm_fn = lambda B, S, K, dev: torch.ones((B, S), dtype=torch.float32, device=dev) + if predict_score_fn is None: + predict_score_fn = lambda B, S, K, dev: torch.zeros( + (B, S, K), dtype=torch.float32, device=dev + ) + if predict_lse_fn is None: + predict_lse_fn = lambda B, S, K, dev: torch.full( + (B, S), float(math.log(max(K, 1))), dtype=torch.float32, device=dev + ) + + fake_dsa = MagicMock(name='_DSA_full_dense_stub') + + def fake_indexer_forward(q_bshd, k_bshd, w_bsh, ratio): + return {'scores': torch.zeros(b, sq, n_comp, dtype=torch.float32, device=q_bshd.device)} + + fake_dsa.indexer_forward_wrapper.side_effect = fake_indexer_forward + + def fake_filtered_topk(scores_flat, seq_lens, top_k, next_n, return_val): + return { + 'indices': torch.zeros( + scores_flat.shape[0], top_k, dtype=torch.int32, device=scores_flat.device + ) + } + + fake_dsa.indexer_top_k_wrapper.side_effect = fake_filtered_topk + + def fake_dense_indexer_score(q, k, w, qhead_per_kv_head, sm_scale, ratio, **kwargs): + dev = q.device + return { + 'out': predict_score_fn(b, sq, n_comp, dev), + 'denom': predict_lse_fn(b, sq, n_comp, dev), + } + + fake_dsa.dense_indexer_score_recompute_wrapper.side_effect = fake_dense_indexer_score + + def fake_dense_attn_score(q, k, lse, softmax_scale, qhead_per_kv_head, ratio, **kwargs): + dev = q.device + return { + 'out': target_score_fn(b, sq, n_comp, dev), + 'denom': target_l1norm_fn(b, sq, n_comp, dev), + } + + fake_dsa.dense_attn_score_recompute_wrapper.side_effect = fake_dense_attn_score + + def fake_sparse_attn_backward(q, kv, out, dout, lse, attn_sink, topk_idxs, **kwargs): + return { + 'dq': torch.full_like(q, dq_value), + 'dkv': torch.full_like(kv, dkv_value), + 'd_sink': torch.full_like(attn_sink, d_sink_value), + } + + fake_dsa.sparse_attention_backward_wrapper.side_effect = fake_sparse_attn_backward + + def fake_dense_indexer_grad_backward( + q_idx_bshd, + w_bsh, + k_idx_bsd, + attn_score, + attn_l1norm, + index_score, + index_lse, + sm_scale, + loss_coeff, + grad_loss, + ratio, + block_I, + ): + return { + 'd_index_q': torch.full_like(q_idx_bshd, d_index_q_value), + 'd_weights': torch.full_like(w_bsh, d_weights_value), + 'd_index_k': torch.full_like(k_idx_bsd, d_index_k_value), + } + + fake_dsa.dense_indexer_backward_wrapper.side_effect = fake_dense_indexer_grad_backward + + flash_stub = _make_flash_mla_stub(d_v=d) + + dk._DSA = fake_dsa + dk._flash_mla_sparse_fwd = flash_stub + return fake_dsa, flash_stub + + +class TestFusedIndexerSparseAttn: + """End-to-end numerical tests for the Path B autograd Function with all + underlying CUDA kernels mocked. + """ + + # Common shapes shared across the forward tests. + SHAPES = dict(sq=4, b=2, np_=2, d=512, skv=8, n_comp=4, idx_nh=4, idx_hd=64) + + def _make_inputs(self, *, requires_grad=False): + """Build the seven differentiable + one non-differentiable inputs.""" + s = self.SHAPES + win_topk = _get_topk_alignment() - 2 # exercise padding + torch.manual_seed(0) + + def make(*shape, dtype, rg=False): + t = torch.randn(*shape, dtype=dtype, device='cuda') + if requires_grad and rg: + t = t.detach().clone().requires_grad_(True) + return t + + query = make(s['sq'], s['b'], s['np_'], s['d'], dtype=torch.bfloat16, rg=True) + kv_full = make(s['skv'], s['b'], s['d'], dtype=torch.bfloat16, rg=True) + attn_sink = torch.zeros(s['np_'], dtype=torch.float32, device='cuda') + if requires_grad: + attn_sink = attn_sink.detach().clone().requires_grad_(True) + window_idxs = torch.zeros(s['b'], s['sq'], win_topk, dtype=torch.int32, device='cuda') + q_indexer = make(s['sq'], s['b'], s['idx_nh'], s['idx_hd'], dtype=torch.bfloat16, rg=True) + k_indexer = make(s['n_comp'], s['b'], s['idx_hd'], dtype=torch.bfloat16, rg=True) + weights = make(s['sq'], s['b'], s['idx_nh'], dtype=torch.bfloat16, rg=True) + return dict( + query=query, + kv_full=kv_full, + attn_sink=attn_sink, + window_idxs=window_idxs, + q_indexer=q_indexer, + k_indexer=k_indexer, + weights=weights, + ) + + @pytest.mark.parametrize( + "loss_coeff, target_kind, expected", + [ + # KL(target == predict) == 0 → loss == 0 regardless of coeff. + (1.0, 'uniform', 0.0), + # loss_coeff == 0 short-circuits even when target != predict. + (0.0, 'peaked', 0.0), + # target = δ_0, predict = uniform(1/K) → KL = log(K) per row, + # mean over rows = log(K), scaled by coeff = coeff * log(K). + (0.7, 'peaked', 0.7 * math.log(2)), + # Linearity in loss_coeff: doubling the coeff doubles the loss. + (2.0, 'peaked', 2.0 * math.log(2)), + ], + ids=['identical_dists', 'coeff_zero', 'analytical_kl', 'linearity_x2'], + ) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_indexer_loss_formula(self, loss_coeff, target_kind, expected, reset_lazy_kernel_state): + """All four loss-property cases share one fixture: + + * KL is zero when target == predict, + * ``loss_coeff == 0`` short-circuits to zero, + * for ``target = δ_0`` and ``predict = uniform(1/K)`` the per-row + KL is exactly ``log(K)`` so the mean is ``loss_coeff * log(K)``, + * the loss is linear in ``loss_coeff``. + """ + s = self.SHAPES + topk = 2 # = effective_topk = min(indexer_topk, n_comp); appears as K + target_fn = ( + _uniform_dist + if target_kind == 'uniform' + else (lambda B, S, K, dev: _peaked_dist(B, S, K, dev, peak_idx=0)) + ) + + inputs = self._make_inputs() + _install_full_dsa_mock( + b=s['b'], + sq=s['sq'], + np_=s['np_'], + d=s['d'], + n_comp=s['n_comp'], + idx_nh=s['idx_nh'], + predict_fn=_uniform_dist, + target_fn=target_fn, + ) + + _, indexer_loss = fused_indexer_sparse_attn( + **inputs, + indexer_topk=topk, + ratio=4, + softmax_scale=0.5, + loss_coeff=loss_coeff, + sparse_loss=True, + kv_offset=s['skv'] - s['n_comp'], + ) + + assert torch.allclose( + indexer_loss, torch.tensor(expected, device='cuda'), rtol=1e-5, atol=1e-5 + ), f"got {indexer_loss.item()}, expected {expected}" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_sparse_path_fwd_output_bwd_grads_and_topk_clamp(self, reset_lazy_kernel_state): + """Combined coverage for the sparse-loss path's three non-numerical + properties (assertion blocks self-label on failure): + + * (a) ``output`` is exactly the FlashMLA stub's ``out`` reshaped + from ``(sq*b, np_, d_v)`` to ``(sq, b, np_ * d_v)``. + * (b) After backward, each leaf gradient equals the corresponding + mocked kernel output, with q/kv/attn_sink coming from the + sparse-attn bwd kernel and q_indexer/k_indexer/weights coming + from the indexer bwd kernel (BSHD → SBHD permute applied). + * (c) ``indexer_topk > n_comp`` is clamped to ``n_comp`` before the + radix TopK kernel is called. + """ + s = self.SHAPES + + # ---- (a) forward pass-through (no grads needed) ------------------ + inputs = self._make_inputs() + _, flash_stub_a = _install_full_dsa_mock( + b=s['b'], sq=s['sq'], np_=s['np_'], d=s['d'], n_comp=s['n_comp'], idx_nh=s['idx_nh'] + ) + output_a, _ = fused_indexer_sparse_attn( + **inputs, + indexer_topk=2, + ratio=4, + softmax_scale=0.5, + indexer_softmax_scale=0.125, + loss_coeff=0.0, + sparse_loss=True, + kv_offset=s['skv'] - s['n_comp'], + ) + assert output_a.shape == (s['sq'], s['b'], s['np_'] * s['d']), "(a) shape" + assert output_a.dtype == torch.bfloat16, "(a) dtype" + expected_a = flash_stub_a.last_out.reshape(s['sq'], s['b'], s['np_'], s['d']).reshape( + s['sq'], s['b'], s['np_'] * s['d'] + ) + assert torch.equal(output_a, expected_a), "(a) forward value pass-through" + + # ---- (b) backward grad propagation ------------------------------- + dk._DSA = None # fresh mocks + dk._flash_mla_sparse_fwd = None + inputs_b = self._make_inputs(requires_grad=True) + _install_full_dsa_mock( + b=s['b'], + sq=s['sq'], + np_=s['np_'], + d=s['d'], + n_comp=s['n_comp'], + idx_nh=s['idx_nh'], + dq_value=7.0, + dkv_value=-3.0, + d_sink_value=11.0, + d_index_q_value=0.5, + d_weights_value=-0.25, + d_index_k_value=1.5, + ) + output_b, indexer_loss_b = fused_indexer_sparse_attn( + **inputs_b, + indexer_topk=2, + ratio=4, + softmax_scale=0.5, + indexer_softmax_scale=0.125, + loss_coeff=1.0, + sparse_loss=True, + kv_offset=s['skv'] - s['n_comp'], + ) + (output_b.sum() + indexer_loss_b).backward() + for name, value in [ + ('query', 7.0), + ('kv_full', -3.0), + ('attn_sink', 11.0), + ('q_indexer', 0.5), + ('k_indexer', 1.5), + ('weights', -0.25), + ]: + grad = inputs_b[name].grad + assert grad is not None, f"(b) {name}: missing grad" + assert torch.equal(grad, torch.full_like(inputs_b[name], value)), ( + f"(b) {name}: grad does not equal full({value}); " + f"got first elem = {grad.float().flatten()[0].item()}" + ) + + # ---- (c) indexer_topk > n_comp clamp ----------------------------- + dk._DSA = None + dk._flash_mla_sparse_fwd = None + inputs_c = self._make_inputs() + fake_dsa_c, _ = _install_full_dsa_mock( + b=s['b'], sq=s['sq'], np_=s['np_'], d=s['d'], n_comp=s['n_comp'], idx_nh=s['idx_nh'] + ) + fused_indexer_sparse_attn( + **inputs_c, + indexer_topk=999, # > n_comp + ratio=4, + softmax_scale=0.5, + loss_coeff=0.0, + sparse_loss=True, + kv_offset=s['skv'] - s['n_comp'], + ) + topk_call = fake_dsa_c.indexer_top_k_wrapper.call_args + assert ( + topk_call.kwargs['top_k'] == s['n_comp'] + ), f"(c) top_k clamp: got {topk_call.kwargs['top_k']}, expected {s['n_comp']}" + + +# --------------------------------------------------------------------------- +# fused_indexer_sparse_attn — dense path (sparse_loss=False) +# --------------------------------------------------------------------------- + + +class TestDenseFusedIndexerSparseAttn: + """End-to-end tests for the dense-loss branch of Path B with all + underlying CUDA kernels mocked. Mirrors :class:`TestFusedIndexerSparseAttn` + but exercises the ``sparse_loss=False`` code path through + :class:`FusedIndexerSparseAttnFunc`. + """ + + SHAPES = dict(sq=4, b=2, np_=2, d=512, skv=8, n_comp=4, idx_nh=4, idx_hd=64) + + def _make_inputs(self, *, requires_grad=False): + s = self.SHAPES + win_topk = _get_topk_alignment() - 2 # exercise padding + torch.manual_seed(0) + + def make(*shape, dtype, rg=False): + t = torch.randn(*shape, dtype=dtype, device='cuda') + if requires_grad and rg: + t = t.detach().clone().requires_grad_(True) + return t + + query = make(s['sq'], s['b'], s['np_'], s['d'], dtype=torch.bfloat16, rg=True) + kv_full = make(s['skv'], s['b'], s['d'], dtype=torch.bfloat16, rg=True) + attn_sink = torch.zeros(s['np_'], dtype=torch.float32, device='cuda') + if requires_grad: + attn_sink = attn_sink.detach().clone().requires_grad_(True) + window_idxs = torch.zeros(s['b'], s['sq'], win_topk, dtype=torch.int32, device='cuda') + q_indexer = make(s['sq'], s['b'], s['idx_nh'], s['idx_hd'], dtype=torch.bfloat16, rg=True) + k_indexer = make(s['n_comp'], s['b'], s['idx_hd'], dtype=torch.bfloat16, rg=True) + weights = make(s['sq'], s['b'], s['idx_nh'], dtype=torch.bfloat16, rg=True) + return dict( + query=query, + kv_full=kv_full, + attn_sink=attn_sink, + window_idxs=window_idxs, + q_indexer=q_indexer, + k_indexer=k_indexer, + weights=weights, + ) + + @pytest.mark.parametrize( + "loss_coeff, target_kind, expected", + [ + # Identical dists: KL == 0 regardless of coeff. + (1.0, 'uniform', 0.0), + # loss_coeff == 0 short-circuits even when target != predict. + (0.0, 'peaked', 0.0), + # target = δ_0, predict = uniform(1/n_comp) + # per-row KL = log(n_comp); mean = log(n_comp); loss = coeff * log(n_comp). + # n_comp = 4 here. + (0.7, 'peaked', 0.7 * math.log(4)), + (2.0, 'peaked', 2.0 * math.log(4)), + ], + ids=['identical_dists', 'coeff_zero', 'analytical_kl', 'linearity_x2'], + ) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dense_indexer_loss_formula( + self, loss_coeff, target_kind, expected, reset_lazy_kernel_state + ): + """``_kl_loss_from_dense_scores`` is the dense analogue of + ``_kl_loss_from_target_predict``. Verifies the same four KL + properties (zero, coeff-zero short-circuit, analytical formula, + linearity in coeff) over the dense ``(B, S_q, S_k)`` tensors. + + We drive the stub outputs so that: + + * predict = ``softmax(0)`` over S_k = uniform(1/n_comp). This is + encoded as ``index_score = 0`` everywhere, ``index_lse = log(n_comp)``; + ``predict = exp(score - lse) = 1/n_comp``. + * For ``target = uniform``: attn_score = 1/n_comp uniformly, attn_l1norm = 1. + * For ``target = δ_0``: attn_score peaked on slot 0 with sum 1, attn_l1norm = 1. + """ + s = self.SHAPES + + if target_kind == 'uniform': + target_score_fn = lambda B, S, K, dev: torch.full( + (B, S, K), 1.0 / max(K, 1), dtype=torch.float32, device=dev + ) + else: + + def target_score_fn(B, S, K, dev): + t = torch.zeros((B, S, K), dtype=torch.float32, device=dev) + t[..., 0] = 1.0 + return t + + target_l1norm_fn = lambda B, S, K, dev: torch.ones((B, S), dtype=torch.float32, device=dev) + + inputs = self._make_inputs() + _install_full_dsa_mock_dense( + b=s['b'], + sq=s['sq'], + np_=s['np_'], + d=s['d'], + n_comp=s['n_comp'], + idx_nh=s['idx_nh'], + target_score_fn=target_score_fn, + target_l1norm_fn=target_l1norm_fn, + # predict_score_fn / predict_lse_fn defaults give uniform predict. + ) + + _, indexer_loss = fused_indexer_sparse_attn( + **inputs, + indexer_topk=2, + ratio=4, + softmax_scale=0.5, + loss_coeff=loss_coeff, + sparse_loss=False, + kv_offset=s['skv'] - s['n_comp'], + ) + + assert torch.allclose( + indexer_loss, torch.tensor(expected, device='cuda'), rtol=1e-5, atol=1e-5 + ), f"got {indexer_loss.item()}, expected {expected}" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dense_path_fwd_kernel_calls_and_bwd_grads(self, reset_lazy_kernel_state): + """Combined coverage for the dense-loss path's two non-numerical + properties (assertion blocks self-label on failure): + + * (a) The forward invokes ``dense_attn_score_recompute_wrapper`` + (NOT the sparse score kernels) with the right BSHD/4-D shapes + and ``ratio`` / scale args. The indexer predict is derived + directly from ``indexer_forward_wrapper`` scores (no separate + ``dense_indexer_score_recompute_wrapper`` call). + * (b) The forward eagerly invokes ``dense_indexer_backward_wrapper`` + (NOT the sparse one), threads ``ratio`` through, and the + resulting grads land on the right SBHD leaves (BSHD → SBHD + permute applied for the indexer-side grads, scaled by + ``grad_loss`` in the actual backward). + """ + s = self.SHAPES + ratio = 4 + softmax_scale = 0.5 + idx_scale = 0.125 + loss_coeff = 1.0 + + # ---- (a) forward kernel selection + arg shapes ------------------- + inputs_a = self._make_inputs() + fake_dsa_a, _ = _install_full_dsa_mock_dense( + b=s['b'], sq=s['sq'], np_=s['np_'], d=s['d'], n_comp=s['n_comp'], idx_nh=s['idx_nh'] + ) + fused_indexer_sparse_attn( + **inputs_a, + indexer_topk=2, + ratio=ratio, + softmax_scale=softmax_scale, + indexer_softmax_scale=idx_scale, + loss_coeff=loss_coeff, + sparse_loss=False, + kv_offset=s['skv'] - s['n_comp'], + ) + # Indexer predict is derived from indexer_forward_wrapper scores + # (gather + logsumexp), NOT from dense_indexer_score_recompute_wrapper. + fake_dsa_a.dense_indexer_score_recompute_wrapper.assert_not_called() + fake_dsa_a.dense_attn_score_recompute_wrapper.assert_called_once() + fake_dsa_a.sparse_indexer_score_recompute_wrapper.assert_not_called() + fake_dsa_a.sparse_attn_score_recompute_wrapper.assert_not_called() + + attn_call = fake_dsa_a.dense_attn_score_recompute_wrapper.call_args + q_attn, k_attn, lse_arg, sm_arg = attn_call.args + assert q_attn.shape == (s['b'], s['sq'], s['np_'], s['d']), "(a) dense attn score: q shape" + assert k_attn.shape == ( + s['b'], + s['n_comp'], + 1, + s['d'], + ), "(a) dense attn score: k shape (h_kv=1)" + assert lse_arg.shape == (s['b'], s['sq'], s['np_']), "(a) dense attn score: lse shape" + assert sm_arg == softmax_scale, "(a) dense attn score: positional softmax_scale" + assert attn_call.kwargs['qhead_per_kv_head'] == s['np_'] + assert attn_call.kwargs['ratio'] == ratio + + # ---- (b) forward-eager indexer backward + grad propagation -------- + dk._DSA = None + dk._flash_mla_sparse_fwd = None + inputs_b = self._make_inputs(requires_grad=True) + fake_dsa_b, _ = _install_full_dsa_mock_dense( + b=s['b'], + sq=s['sq'], + np_=s['np_'], + d=s['d'], + n_comp=s['n_comp'], + idx_nh=s['idx_nh'], + dq_value=7.0, + dkv_value=-3.0, + d_sink_value=11.0, + d_index_q_value=0.5, + d_weights_value=-0.25, + d_index_k_value=1.5, + ) + output, indexer_loss = fused_indexer_sparse_attn( + **inputs_b, + indexer_topk=2, + ratio=ratio, + softmax_scale=softmax_scale, + indexer_softmax_scale=idx_scale, + loss_coeff=loss_coeff, + sparse_loss=False, + kv_offset=s['skv'] - s['n_comp'], + ) + (output.sum() + indexer_loss).backward() + + # dense_indexer_backward_wrapper is called eagerly during forward. + fake_dsa_b.dense_indexer_backward_wrapper.assert_called_once() + fake_dsa_b.indexer_backward_wrapper.assert_not_called() + ig_call = fake_dsa_b.dense_indexer_backward_wrapper.call_args + assert ig_call.kwargs['ratio'] == ratio, "(b) ratio not threaded through" + assert ig_call.kwargs['sm_scale'] == idx_scale, "(b) sm_scale not threaded" + assert ig_call.kwargs['loss_coeff'] == loss_coeff, "(b) loss_coeff not threaded" + + for name, value in [ + ('query', 7.0), + ('kv_full', -3.0), + ('attn_sink', 11.0), + ('q_indexer', 0.5), + ('k_indexer', 1.5), + ('weights', -0.25), + ]: + grad = inputs_b[name].grad + assert grad is not None, f"(b) {name}: missing grad" + assert torch.equal( + grad, torch.full_like(inputs_b[name], value) + ), f"(b) {name}: grad does not equal full({value})" + + +# --------------------------------------------------------------------------- +# Real-kernel parity tests (cuDNN + optional FlashMLA) +# --------------------------------------------------------------------------- +# +# Everything above this banner stubs ``cudnn.DSA`` and ``flash_mla`` with +# ``MagicMock``-based fakes; that exercises the Python plumbing of +# ``csa_kernels.py`` (shape transforms, autograd wiring, KL composition) +# but does NOT verify that the cuDNN kernels themselves compute what +# ``csa_kernels.py`` expects them to compute. +# +# The tests below close that gap by running each helper / public function +# end-to-end against a small PyTorch reference implementation. Numeric +# tolerances are bf16-friendly (atol/rtol ~ 5e-2 for raw scores, 1e-3 for +# normalized distributions, 5e-2 for backward grads). +# +# Skipped automatically when: +# * CUDA is unavailable; +# * cuDNN frontend is not installed (``import cudnn`` fails); +# * ``cudnn.DSA`` namespace is missing; +# * SM is too low (sparse: SM90+; dense: SM100+); +# * for FlashMLA-needing tests, ``flash_mla`` is not installed. +# --------------------------------------------------------------------------- + + +def _skip_if_real_kernels_unavailable(*, sm_min: int = 9, need_flash_mla: bool = False): + """Pytest-side gate for real-kernel tests. Raises ``pytest.skip`` if + any of the runtime dependencies are missing. + """ + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + sm_major = torch.cuda.get_device_capability()[0] + if sm_major < sm_min: + pytest.skip(f"requires SM{sm_min}+, found SM{sm_major}") + cudnn = pytest.importorskip("cudnn") + cudnn_frontend = pytest.importorskip("cudnn_frontend") + from packaging.version import Version + + if Version(cudnn_frontend.__version__) < Version("1.24.0"): + pytest.skip(f"requires cudnn_frontend>=1.24.0, found {cudnn_frontend.__version__}") + if not hasattr(cudnn, 'DSA'): + pytest.skip("cudnn.DSA namespace not available") + if need_flash_mla: + pytest.importorskip("flash_mla") + + +# --------------------------------------------------------------------------- +# PyTorch reference implementations +# --------------------------------------------------------------------------- + + +def _ratio_causal_valid_mask(sq: int, sk: int, ratio: int, device) -> torch.Tensor: + """``(Sq, Sk)`` bool: valid iff ``k_idx < min(Sk, (q_idx + 1) // ratio)``. + + Matches the cuDNN dense-score kernels' built-in causal mask + (``col_limit = min(S_k, (q + 1) // ratio)``) and ``csa.py``'s + ``compress_ratio`` mask formulation. + """ + q_idx = torch.arange(sq, device=device).unsqueeze(1) # (Sq, 1) + k_idx = torch.arange(sk, device=device).unsqueeze(0) # (1, Sk) + col_limit = ((q_idx + 1) // ratio).clamp(max=sk) + return k_idx < col_limit # (Sq, Sk) + + +def _ref_indexer_full_score( + q_bshd_fp32: torch.Tensor, # (B, Sq, H, D) + k_bsd_fp32: torch.Tensor, # (B, Sk, D) — MQA + w_bsh_fp32: torch.Tensor, # (B, Sq, H) + sm_scale: float, + ratio: int, +) -> torch.Tensor: + """Reference for ``_bwd_dense_indexer_score.out``. + + ``S[b,q,k] = sm_scale * sum_h ReLU(Q[b,q,h] @ K[b,k]^T) * W[b,q,h]``, + with the kernel's bottom-right ratio causal mask producing ``-inf`` + at masked positions. + """ + B, Sq, _, _ = q_bshd_fp32.shape + Sk = k_bsd_fp32.shape[1] + qk = torch.einsum('bqhd,bkd->bqhk', q_bshd_fp32, k_bsd_fp32) # (B, Sq, H, Sk) + relu_qk = torch.relu(qk) + s = (relu_qk * w_bsh_fp32.unsqueeze(-1)).sum(dim=2) * sm_scale # (B, Sq, Sk) + valid = _ratio_causal_valid_mask(Sq, Sk, ratio, s.device).unsqueeze(0) + return torch.where(valid, s, torch.full_like(s, float('-inf'))) + + +def _ref_attn_full_score( + q_bshd_fp32: torch.Tensor, # (B, Sq, H, D) + k_bsd_fp32: torch.Tensor, # (B, Sk, D) — MQA + lse_bshq_fp32: torch.Tensor, # (B, Sq, H) + softmax_scale: float, + ratio: int, +) -> torch.Tensor: + """Reference for ``_bwd_dense_attn_score.out``. + + ``out[b,q,k] = sum_h exp(Q[b,q,h] @ K[b,k]^T * scale - LSE[b,q,h])``, + with the ratio causal mask producing ``0`` at masked positions + (the per-head ``exp`` is zeroed out, contributing nothing to the sum). + """ + B, Sq, _, _ = q_bshd_fp32.shape + Sk = k_bsd_fp32.shape[1] + qk = torch.einsum('bqhd,bkd->bqhk', q_bshd_fp32, k_bsd_fp32) * softmax_scale + p = torch.exp(qk - lse_bshq_fp32.unsqueeze(-1)) + s = p.sum(dim=2) # (B, Sq, Sk) + valid = _ratio_causal_valid_mask(Sq, Sk, ratio, s.device).unsqueeze(0) + return torch.where(valid, s, torch.zeros_like(s)) + + +def _ref_indexer_predict_sparse(q_bshd_fp32, k_bsd_fp32, w_bsh_fp32, topk_indices, sm_scale): + """Reference for ``sparse_indexer_score_recompute_wrapper.predict``. + + Compute the full-KV indexer score, gather ``topk_indices``, softmax + over the topK axis. ``-1`` entries in topk are masked to ``-inf`` + so they contribute zero probability. + """ + qk = torch.einsum('bqhd,bkd->bqhk', q_bshd_fp32, k_bsd_fp32) + s = (torch.relu(qk) * w_bsh_fp32.unsqueeze(-1)).sum(dim=2) * sm_scale # (B, Sq, Sk) + valid = topk_indices >= 0 + safe = topk_indices.clamp(min=0).long() + s_topk = torch.gather(s, dim=-1, index=safe) + s_topk = torch.where(valid, s_topk, torch.full_like(s_topk, float('-inf'))) + return torch.softmax(s_topk, dim=-1) + + +def _ref_attn_target_sparse(q_bshd_fp32, k_bsd_fp32, lse_bsh_fp32, topk_indices, softmax_scale): + """Reference for ``sparse_attn_score_recompute_wrapper.target``. + + Per-head ``exp(QK*scale - LSE)``, sum over heads, gather topK, + L1-normalise over the topK axis. ``-1`` entries are zero-masked + pre-normalisation. + """ + qk = torch.einsum('bqhd,bkd->bqhk', q_bshd_fp32, k_bsd_fp32) * softmax_scale + p = torch.exp(qk - lse_bsh_fp32.unsqueeze(-1)) # (B, Sq, H, Sk) + s = p.sum(dim=2) # (B, Sq, Sk) + valid = topk_indices >= 0 + safe = topk_indices.clamp(min=0).long() + s_topk = torch.gather(s, dim=-1, index=safe) + s_topk = torch.where(valid, s_topk, torch.zeros_like(s_topk)) + denom = s_topk.sum(dim=-1, keepdim=True).clamp(min=1e-10) + return s_topk / denom + + +def _ref_dense_indexer_loss( + q_indexer_bshd_fp32, + k_indexer_bsd_fp32, + w_bsh_fp32, + q_attn_bshd_fp32, + k_attn_bsd_fp32, + lse_bshq_fp32, + indexer_softmax_scale: float, + attn_softmax_scale: float, + ratio: int, + loss_coeff: float, +) -> torch.Tensor: + """Reference dense KL loss (matches ``compute_dsa_indexer_loss(sparse_loss=False)`` + in ``dsa.py``). Uses the same ratio causal mask the kernel applies. + """ + eps = torch.finfo(torch.float32).tiny + # Per-(b,q,k) raw scores via the same formulas the kernels use. + attn_scores = _ref_attn_full_score( + q_attn_bshd_fp32, k_attn_bsd_fp32, lse_bshq_fp32, attn_softmax_scale, ratio + ) # (B, Sq, Sk) head-summed, ratio-masked, zeros at masked positions. + index_scores = _ref_indexer_full_score( + q_indexer_bshd_fp32, k_indexer_bsd_fp32, w_bsh_fp32, indexer_softmax_scale, ratio + ) # (B, Sq, Sk) ReLU·W, ratio-masked, -inf at masked positions. + + # L1-norm denom for target; LSE for predict. + attn_denom = attn_scores.sum(dim=-1) # (B, Sq) + index_lse = torch.logsumexp(index_scores, dim=-1) # (B, Sq), -inf for fully-masked rows + + row_valid = (attn_denom > eps) & torch.isfinite(index_lse) + + safe_l1 = attn_denom.clamp(min=eps) + safe_lse = torch.where(row_valid, index_lse, torch.zeros_like(index_lse)) + + target = attn_scores / safe_l1.unsqueeze(-1) + target_clamped = target.clamp(min=eps) + # Mask within-row: ratio-causal-masked positions have ``index_scores = + # -inf`` (from ``_ref_indexer_full_score``). Letting them flow into + # ``log_predict`` would make per-position contributions blow up to + # +inf (``target_clamped * (log(target) - (-inf)) = +inf``). They have + # zero mass under ``target`` (``_ref_attn_full_score`` zeros those + # positions) so their KL contribution should be 0; explicitly mask. + position_valid = torch.isfinite(index_scores) + log_predict = torch.where( + position_valid, index_scores - safe_lse.unsqueeze(-1), torch.zeros_like(index_scores) + ) + contributions = target_clamped * (torch.log(target_clamped) - log_predict) + contributions = torch.where(position_valid, contributions, torch.zeros_like(contributions)) + kl_per_row = contributions.sum(dim=-1) + kl_per_row = torch.where(row_valid, kl_per_row, torch.zeros_like(kl_per_row)) + return loss_coeff * kl_per_row.mean() + + +def _ref_sparse_attn_forward( + q_flat_bf16: torch.Tensor, # (total_Sq, H, D) + kv_flat_bf16: torch.Tensor, # (total_Skv, D) — K=V, MQA + attn_sink_fp32: torch.Tensor, # (H,) + topk_idxs: torch.Tensor, # (total_Sq, topk) int32, global + softmax_scale: float, + d_v: int, +): + """Pure-PyTorch reference for FlashMLA sparse-attn-fwd output. + + Mirrors the math FlashMLA implements: + * Scores ``S[i, h, k] = Q[i, h] @ K[topk[i, k]]^T * scale`` for valid ``k``; + * Append a per-head sink logit (``attn_sink``); + * ``softmax`` over the (topk + sink) axis; + * ``out[i, h] = sum_k softmax[i, h, k] * V[topk[i, k]]`` (excluding sink). + + Returns ``(out, lse)`` in the same shapes/dtype as the FlashMLA kernel. + Invalid ``-1`` topk entries contribute zero to the softmax (logit -inf). + """ + total_Sq, H, D = q_flat_bf16.shape + topk = topk_idxs.shape[-1] + device = q_flat_bf16.device + q_fp32 = q_flat_bf16.float() + kv_fp32 = kv_flat_bf16.float() + + valid = topk_idxs >= 0 # (total_Sq, topk) + safe = topk_idxs.clamp(min=0).long() + k_gathered = kv_fp32[safe] # (total_Sq, topk, D) + + qk = torch.einsum('ihd,ikd->ihk', q_fp32, k_gathered) * softmax_scale # (Sq, H, topk) + qk = torch.where(valid.unsqueeze(1).expand(-1, H, -1), qk, torch.full_like(qk, float('-inf'))) + sink = attn_sink_fp32.view(1, H, 1).expand(total_Sq, H, 1) # logit + logits = torch.cat([qk, sink], dim=-1) # (Sq, H, topk + 1) + probs = torch.softmax(logits, dim=-1) # numerically stable + probs_kv = probs[..., :topk] # exclude sink contribution from output + + v_gathered = k_gathered # K = V (MQA, head-broadcast) + out_fp32 = torch.einsum('ihk,ikd->ihd', probs_kv, v_gathered) # (Sq, H, D_v=D) + if d_v != D: + out_fp32 = out_fp32[..., :d_v] + out = out_fp32.to(q_flat_bf16.dtype) + + # FlashMLA's KV-only LSE excludes the sink term: + # lse_kv[i, h] = logsumexp_k(qk[i, h, k]) over valid k only. + lse_kv = torch.logsumexp(qk, dim=-1) # (Sq, H), -inf for fully-masked rows + return out, lse_kv + + +# --------------------------------------------------------------------------- +# Score-helper parity tests (sparse + dense): real cuDNN vs PyTorch reference +# --------------------------------------------------------------------------- + + +# Shared small shape across all real-kernel tests to maximize cuDNN compile-cache +# hits. ``ratio=1`` (standard upper-triangular causal) keeps the math simple +# and ensures every row has at least one valid KV position. +_REAL_SHAPES_SPARSE = dict( + b=2, + sq=128, + sk=128, + n_comp=128, + np_=32, + 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=128, + ratio=1, + softmax_scale=512**-0.5, + indexer_softmax_scale=128**-0.5, +) + + +def _build_real_score_inputs(s, *, with_lse: bool = True, with_topk: bool = True): + """Build a coherent set of bf16 BSHD inputs for the score-helper tests. + + Returns a dict with both bf16 (kernel-ready) and fp32 (reference-math) + views of every tensor, plus optional LSE / topk_indices. + """ + torch.manual_seed(0) + dev = 'cuda' + + q_idx = torch.randn(s['b'], s['sq'], s['idx_nh'], s['idx_hd'], dtype=torch.bfloat16, device=dev) + k_idx = torch.randn(s['b'], s['sk'], s['idx_hd'], dtype=torch.bfloat16, device=dev) + w = torch.randn(s['b'], s['sq'], s['idx_nh'], dtype=torch.bfloat16, device=dev) + + q_attn = torch.randn(s['b'], s['sq'], s['np_'], s['d'], dtype=torch.bfloat16, device=dev) + k_attn = torch.randn(s['b'], s['sk'], s['d'], dtype=torch.bfloat16, device=dev) + + out = dict(q_idx=q_idx, k_idx=k_idx, w=w, q_attn=q_attn, k_attn=k_attn) + + if with_lse: + # LSE = logsumexp(QK*scale, dim=Sk) with the kernel's ratio mask. + # Real LSE input avoids exp(-inf - finite) underflow during reference. + qk = torch.einsum('bqhd,bkd->bqhk', q_attn.float(), k_attn.float()) * s['softmax_scale'] + valid = _ratio_causal_valid_mask(s['sq'], s['sk'], s['ratio'], qk.device).view( + 1, s['sq'], 1, s['sk'] + ) + qk_masked = torch.where(valid, qk, torch.full_like(qk, float('-inf'))) + out['lse'] = torch.logsumexp(qk_masked, dim=-1).clamp(min=-1e30).contiguous() + + if with_topk: + # Pick distinct random valid indices per (b, sq), with a few -1s + # interleaved to exercise the invalid-slot path. + topk = s['topk'] + torch.manual_seed(123) + idxs = torch.randint(0, s['sk'], (s['b'], s['sq'], topk), dtype=torch.int32, device=dev) + # Mark a few slots invalid (-1) to test the topk_indices < 0 path. + invalid = torch.rand(s['b'], s['sq'], topk, device=dev) < 0.1 + idxs = torch.where(invalid, torch.full_like(idxs, -1), idxs) + out['topk'] = idxs + + return out + + +class TestRealKernelScoreHelpers: + """Real-kernel parity tests for the four ``_compute_*`` score helpers + against PyTorch reference implementations. Single parametrized test + covers all four; numeric tolerance is bf16-friendly (raw fp32 score + sums agree to ~5%, normalized distributions to ~5e-3). + """ + + # Each case: (id, sm_min, kernel_name, runner). The runner does the + # call + ref + assertion; it returns nothing on success. + @pytest.mark.parametrize( + "case", + ['sparse_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) + + s = _REAL_SHAPES_SPARSE + # Each case needs a different combination of the input fixture. + x = _build_real_score_inputs( + s, + with_lse=case.endswith('_attn_target') or case.endswith('_attn_score'), + with_topk=case.startswith('sparse_'), + ) + + from megatron.core.transformer.experimental_attention_variant import csa_kernels as _dk + + if case == 'sparse_indexer_predict': + # The kernel takes sm_scale=1.0; scale is applied via weights + # pre-multiplication (relu(c·x)·W trick). Reference mirrors that. + scale = s['indexer_softmax_scale'] + w_scaled = (x['w'].float() * scale).to(x['w'].dtype) + out = _dk._compute_indexer_predict( + x['q_idx'], x['k_idx'], w_scaled, x['topk'], qhead_per_kv_head=s['idx_nh'] + ) + ref = _ref_indexer_predict_sparse( + x['q_idx'].float(), x['k_idx'].float(), w_scaled.float(), x['topk'], sm_scale=1.0 + ) + # Softmax outputs in [0, 1]; bf16 element-wise noise can break + # absolute tolerance, so compare directions via cosine similarity. + assert out.shape == ref.shape == (s['b'], s['sq'], s['topk']) + cos = torch.nn.functional.cosine_similarity( + out.flatten().unsqueeze(0).float(), ref.flatten().unsqueeze(0).float() + ).item() + assert cos > 0.99, ( + f"{case}: cos sim = {cos:.4f}, " + f"max abs diff = {(out - ref).abs().max().item():.3e}" + ) + + elif case == 'sparse_attn_target': + out = _dk._compute_attn_target( + x['q_attn'], + x['k_attn'], + x['lse'], + x['topk'], + softmax_scale=s['softmax_scale'], + qhead_per_kv_head=s['np_'], + ) + ref = _ref_attn_target_sparse( + x['q_attn'].float(), + x['k_attn'].float(), + x['lse'], + x['topk'], + softmax_scale=s['softmax_scale'], + ) + assert out.shape == ref.shape == (s['b'], s['sq'], s['topk']) + cos = torch.nn.functional.cosine_similarity( + out.flatten().unsqueeze(0).float(), ref.flatten().unsqueeze(0).float() + ).item() + assert cos > 0.99, ( + f"{case}: cos sim = {cos:.4f}, " + f"max abs diff = {(out - ref).abs().max().item():.3e}" + ) + + elif case == 'dense_indexer_score': + out, denom = _dk._compute_dense_indexer_score( + x['q_idx'], + x['k_idx'].unsqueeze(2), + x['w'], + qhead_per_kv_head=s['idx_nh'], + indexer_softmax_scale=s['indexer_softmax_scale'], + ratio=s['ratio'], + ) + ref_out = _ref_indexer_full_score( + x['q_idx'].float(), + x['k_idx'].float(), + x['w'].float(), + sm_scale=s['indexer_softmax_scale'], + ratio=s['ratio'], + ) + ref_denom = torch.logsumexp(ref_out, dim=-1) + assert out.shape == ref_out.shape == (s['b'], s['sq'], s['sk']) + assert denom.shape == ref_denom.shape == (s['b'], s['sq']) + # Raw fp32 score sums: relative tolerance dominates. Compare + # only valid positions (masked = -inf in both, NaN under sub). + valid = ( + _ratio_causal_valid_mask(s['sq'], s['sk'], s['ratio'], out.device) + .unsqueeze(0) + .expand_as(out) + ) + diff = torch.where(valid, (out - ref_out).abs(), torch.zeros_like(out)) + scale = ref_out.where(valid, torch.zeros_like(ref_out)).abs().max().item() + assert diff.max().item() <= max( + 5e-2, 5e-2 * scale + ), f"{case}: max abs diff = {diff.max().item():.3e}, scale = {scale:.3e}" + row_valid = torch.isfinite(ref_denom) + assert torch.allclose( + denom[row_valid], ref_denom[row_valid], atol=5e-3, rtol=5e-2 + ), f"{case}: LSE max abs diff = {(denom - ref_denom)[row_valid].abs().max().item():.3e}" + + elif case == 'dense_attn_score': + out, denom = _dk._compute_dense_attn_score( + x['q_attn'], + x['k_attn'].unsqueeze(2), + x['lse'], + qhead_per_kv_head=s['np_'], + softmax_scale=s['softmax_scale'], + ratio=s['ratio'], + ) + ref_out = _ref_attn_full_score( + x['q_attn'].float(), + x['k_attn'].float(), + x['lse'], + softmax_scale=s['softmax_scale'], + ratio=s['ratio'], + ) + ref_denom = ref_out.sum(dim=-1) + assert out.shape == ref_out.shape == (s['b'], s['sq'], s['sk']) + assert denom.shape == ref_denom.shape == (s['b'], s['sq']) + valid = ( + _ratio_causal_valid_mask(s['sq'], s['sk'], s['ratio'], out.device) + .unsqueeze(0) + .expand_as(out) + ) + diff = torch.where(valid, (out - ref_out).abs(), torch.zeros_like(out)) + # exp(QK*scale - LSE) outputs in (0, ~1]: absolute dominates. + assert diff.max().item() <= 5e-3, f"{case}: max abs diff = {diff.max().item():.3e}" + assert torch.allclose(denom, ref_denom, atol=5e-3, rtol=5e-2), ( + f"{case}: denom max abs diff = " f"{(denom - ref_denom).abs().max().item():.3e}" + ) + + else: + raise AssertionError(f"unknown case: {case}") + + +# --------------------------------------------------------------------------- +# KL loss reference parity (dense path; sparse already CPU-tested above). +# --------------------------------------------------------------------------- + + +class TestRealKernelKLLossDense: + """End-to-end parity for ``_kl_loss_from_dense_scores``: run the real + cuDNN dense score kernels, feed their outputs into the helper, and + compare the KL value to the all-PyTorch reference. + """ + + @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) + from megatron.core.transformer.experimental_attention_variant.csa_kernels import ( + _compute_dense_attn_score, + _compute_dense_indexer_score, + _kl_loss_from_dense_scores, + ) + + s = _REAL_SHAPES_SPARSE + x = _build_real_score_inputs(s, with_lse=True, with_topk=False) + loss_coeff = 0.5 + + index_score, index_lse = _compute_dense_indexer_score( + x['q_idx'], + x['k_idx'].unsqueeze(2), + x['w'], + qhead_per_kv_head=s['idx_nh'], + indexer_softmax_scale=s['indexer_softmax_scale'], + ratio=s['ratio'], + ) + attn_score, attn_l1norm = _compute_dense_attn_score( + x['q_attn'], + x['k_attn'].unsqueeze(2), + x['lse'], + qhead_per_kv_head=s['np_'], + softmax_scale=s['softmax_scale'], + ratio=s['ratio'], + ) + + loss_actual = _kl_loss_from_dense_scores( + attn_score, attn_l1norm, index_score, index_lse, loss_coeff + ) + loss_ref = _ref_dense_indexer_loss( + x['q_idx'].float(), + x['k_idx'].float(), + x['w'].float(), + x['q_attn'].float(), + x['k_attn'].float(), + x['lse'], + indexer_softmax_scale=s['indexer_softmax_scale'], + attn_softmax_scale=s['softmax_scale'], + ratio=s['ratio'], + loss_coeff=loss_coeff, + ) + assert torch.allclose(loss_actual, loss_ref, atol=1e-3, rtol=1e-2), ( + f"actual = {loss_actual.item():.6f}, ref = {loss_ref.item():.6f}, " + f"abs diff = {(loss_actual - loss_ref).abs().item():.3e}" + ) + + +# --------------------------------------------------------------------------- +# Real ``indexer_topk``: the top-K set should match the reference ranking. +# --------------------------------------------------------------------------- + + +class TestRealKernelIndexerTopk: + """Real-kernel parity for :func:`indexer_topk`: the SET of selected + top-K indices must match a PyTorch reference ranking. We compare sets + rather than ordered lists because BF16 ties may be broken differently. + """ + + @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+ + from megatron.core.transformer.experimental_attention_variant.csa_kernels import ( + indexer_topk, + ) + + # IndexerForward requires idx_hd=128 and qhpkv in (32, 64). Use an + # SBHD shape that matches what csa.py produces (tensors are SBHD, + # ratio is the indexer's compression ratio). b=2 exercises the + # batch-aware ``seq_lens.repeat(b)`` and the ``(b*sq, sk) → (b, sq, + # topk)`` reshape inside ``_indexer_topk_core`` (BSHD branch). + s = dict( + b=2, + sq=128, + idx_nh=32, + idx_hd=128, + sk=128, + indexer_topk=8, + ratio=4, + indexer_softmax_scale=128**-0.5, + ) + torch.manual_seed(0) + dev = 'cuda' + q_indexer = torch.randn( + s['sq'], s['b'], s['idx_nh'], s['idx_hd'], dtype=torch.bfloat16, device=dev + ) + k_indexer = torch.randn(s['sk'], s['b'], s['idx_hd'], dtype=torch.bfloat16, device=dev) + weights = torch.randn(s['sq'], s['b'], s['idx_nh'], dtype=torch.bfloat16, device=dev) + + topk_indices, topk_length = indexer_topk( + q_indexer, + k_indexer, + weights, + topk=s['indexer_topk'], + ratio=s['ratio'], + indexer_softmax_scale=s['indexer_softmax_scale'], + ) + assert topk_indices.shape == (s['b'], s['sq'], s['indexer_topk']) + assert topk_indices.dtype == torch.int32 + + # Reference: full indexer score, ratio causal mask, take top-K per row + # by descending score. Score is sm_scale * sum_h ReLU(Q@K) * W. + q_bshd = q_indexer.permute(1, 0, 2, 3).contiguous().float() + k_bsd = k_indexer.permute(1, 0, 2).contiguous().float() + w_bsh = weights.permute(1, 0, 2).contiguous().float() + ref_scores = _ref_indexer_full_score( + q_bshd, k_bsd, w_bsh, sm_scale=s['indexer_softmax_scale'], ratio=s['ratio'] + ) # (B, Sq, Sk), -inf at masked positions + + # For each row, count valid positions (un-masked). topk_length should + # match min(indexer_topk, num_valid). + n_valid = (ref_scores > float('-inf')).sum(dim=-1) # (B, Sq) + expected_length = n_valid.clamp(max=s['indexer_topk']).int() + assert torch.equal(topk_length, expected_length) + + # Set comparison row-by-row. Skip rows with 0 valid (kernel returns + # all -1; reference picks arbitrary -inf positions). + ref_topk = torch.topk(ref_scores, k=s['indexer_topk'], dim=-1).indices + for bi in range(s['b']): + for qi in range(s['sq']): + n = int(expected_length[bi, qi].item()) + if n == 0: + # Kernel must report all -1. + assert torch.all(topk_indices[bi, qi] == -1) + continue + actual_set = set(topk_indices[bi, qi, :n].tolist()) + ref_set = set(ref_topk[bi, qi, :n].tolist()) + # BF16 ties may differ: allow up to ~10% mismatch on small K. + inter = actual_set & ref_set + assert len(inter) >= max(1, n - 1), ( + f"row (b={bi}, q={qi}): " + f"actual {sorted(actual_set)} vs ref {sorted(ref_set)}" + ) + + +# --------------------------------------------------------------------------- +# Real ``dsa_sparse_attn``: forward + backward parity vs PyTorch reference. +# --------------------------------------------------------------------------- + + +class TestRealKernelDsaSparseAttn: + """Real-kernel parity for :func:`dsa_sparse_attn`. Forward uses real + FlashMLA + the SBHD/flat reshape wrapper; backward uses the real cuDNN + sparse-attn-bwd kernel. Both checked in one test against the + pure-PyTorch sparse-attn reference (``_ref_sparse_attn_forward``). + """ + + SHAPES = dict(b=2, sq=128, np_=64, d=512, skv=128, topk=32, softmax_scale=512**-0.5) + + def _make_inputs(self, *, requires_grad: bool): + s = self.SHAPES + torch.manual_seed(0) + dev = 'cuda' + + def make_leaf(*shape, dtype): + t = torch.randn(*shape, dtype=dtype, device=dev) + return t.detach().clone().requires_grad_(True) if requires_grad else t + + query = make_leaf(s['sq'], s['b'], s['np_'], s['d'], dtype=torch.bfloat16) + kv = make_leaf(s['skv'], s['b'], s['d'], dtype=torch.bfloat16) + attn_sink = torch.zeros(s['np_'], dtype=torch.float32, device=dev) + if requires_grad: + attn_sink = attn_sink.detach().clone().requires_grad_(True) + + # Coherent valid global topk indices in SBHD-flat layout, with a + # standard causal mask (index <= q_idx). + torch.manual_seed(1) + topk_local = torch.randint( + 0, s['skv'], (s['b'], s['sq'], s['topk']), dtype=torch.int64, device=dev + ) + q_idx = torch.arange(s['sq'], device=dev).view(1, -1, 1) + topk_local = torch.minimum(topk_local, q_idx) + global_idxs = local_to_global_flat(topk_local, s['b']).contiguous() + return query, kv, attn_sink, global_idxs + + def test_real_dsa_sparse_attn_fwd_bwd_matches_reference(self, reset_lazy_kernel_state): + """Forward output AND backward gradients (dq, dkv, d_sink) must + match a pure-PyTorch sparse-attn reference. Combining both checks + 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) + s = self.SHAPES + + # ---- Real path: forward + backward via dsa_sparse_attn ---- + query, kv, attn_sink, global_idxs = self._make_inputs(requires_grad=True) + out = dsa_sparse_attn(query, kv, attn_sink, global_idxs, softmax_scale=s['softmax_scale']) + torch.manual_seed(7) + upstream = torch.randn_like(out) + (out * upstream).sum().backward() + dq_actual = query.grad.float().clone() + dkv_actual = kv.grad.float().clone() + dsink_actual = attn_sink.grad.float().clone() + out_actual = out.float().detach().clone() + + # ---- Reference: pure-PyTorch forward + autograd ---- + query_ref, kv_ref, attn_sink_ref, _ = self._make_inputs(requires_grad=True) + q_flat = query_ref.reshape(s['sq'] * s['b'], s['np_'], s['d']) + kv_flat = kv_ref.reshape(s['skv'] * s['b'], s['d']) + ref_out_flat, _ = _ref_sparse_attn_forward( + q_flat, + kv_flat, + attn_sink_ref, + global_idxs, + softmax_scale=s['softmax_scale'], + d_v=s['d'], + ) + ref_out = ref_out_flat.reshape(s['sq'], s['b'], s['np_'], s['d']).reshape( + s['sq'], s['b'], s['np_'] * s['d'] + ) + (ref_out * upstream).sum().backward() + + # ---- Forward + backward parity (cos sim) ---- + # bf16 GEMM accumulators in FlashMLA fwd / cuDNN sparse-attn-bwd + # make element-wise tolerances brittle (esp. dkv); compare each + # tensor's direction via cosine similarity instead. + def _cos(a, b): + return torch.nn.functional.cosine_similarity( + a.flatten().unsqueeze(0).float(), b.flatten().unsqueeze(0).float() + ).item() + + assert out_actual.shape == ref_out.shape + for name, actual, ref in [ + ('forward', out_actual, ref_out.float()), + ('dq', dq_actual, query_ref.grad.float()), + ('dkv', dkv_actual, kv_ref.grad.float()), + ('d_sink', dsink_actual, attn_sink_ref.grad.float()), + ]: + cos = _cos(actual, ref) + assert cos > 0.99, ( + f"{name}: cos sim = {cos:.4f}, " + f"max abs diff = {(actual - ref).abs().max().item():.3e}" + ) + + +# --------------------------------------------------------------------------- +# Real ``fused_indexer_sparse_attn``: dense-loss path end-to-end parity. +# --------------------------------------------------------------------------- + + +class TestRealKernelFusedIndexerSparseAttn: + """End-to-end parity for the dense loss path of + :func:`fused_indexer_sparse_attn`: real cuDNN dense kernels (forward + + backward) + real FlashMLA, compared to ``_ref_dense_indexer_loss``. + + Backward grad correctness for the indexer-grad kernel is established by + ``TestRealKernelKLLossDense`` (kernel-level math) and + ``TestDenseFusedIndexerSparseAttn::test_dense_backward_calls_dense_indexer_grad`` + (mock-based plumbing). This class only checks the loss SCALAR value. + """ + + # FlashMLA only accepts indexer_topk ∈ {0, 512, 1024, 2048} and a limited + # set of h_q values (np_=64 is the supported one used by the sibling + # DsaSparseAttn real-kernel test). n_comp must be ≥ indexer_topk; skv ≥ + # n_comp so kv_offset = skv - n_comp > 0 still exercises the offset path. + SHAPES = dict( + b=2, + sq=128, + np_=64, + d=512, + skv=640, + n_comp=512, + idx_nh=32, + idx_hd=128, + indexer_topk=512, + ratio=4, + win_topk=8, + softmax_scale=512**-0.5, + indexer_softmax_scale=128**-0.5, + ) + + def test_real_fused_dense_loss_matches_reference(self, reset_lazy_kernel_state): + """Real dense path's KL loss value matches the all-PyTorch reference + on the same inputs. The reference uses an analytical + ``logsumexp(QK*scale, ratio mask)`` for ``lse_indexer`` (FlashMLA + 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) + s = self.SHAPES + torch.manual_seed(0) + dev = 'cuda' + loss_coeff = 0.5 + + # Build inputs once; share between actual and reference. + query = torch.randn(s['sq'], s['b'], s['np_'], s['d'], dtype=torch.bfloat16, device=dev) + kv_full = torch.randn(s['skv'], s['b'], s['d'], dtype=torch.bfloat16, device=dev) + attn_sink = torch.zeros(s['np_'], dtype=torch.float32, device=dev) + torch.manual_seed(1) + win_idxs = torch.randint( + 0, s['sq'], (s['b'], s['sq'], s['win_topk']), dtype=torch.int32, device=dev + ) + q_indexer = torch.randn( + s['sq'], s['b'], s['idx_nh'], s['idx_hd'], dtype=torch.bfloat16, device=dev + ) + k_indexer = torch.randn(s['n_comp'], s['b'], s['idx_hd'], dtype=torch.bfloat16, device=dev) + weights = torch.randn(s['sq'], s['b'], s['idx_nh'], dtype=torch.bfloat16, device=dev) + kv_offset = s['skv'] - s['n_comp'] + + # Real path. + _, indexer_loss = fused_indexer_sparse_attn( + query, + kv_full, + attn_sink, + win_idxs, + q_indexer, + k_indexer, + weights, + indexer_topk=s['indexer_topk'], + ratio=s['ratio'], + softmax_scale=s['softmax_scale'], + indexer_softmax_scale=s['indexer_softmax_scale'], + loss_coeff=loss_coeff, + sparse_loss=False, + kv_offset=kv_offset, + ) + + # Reference: SBHD->BSHD once, build analytical lse_ref, compute KL. + q_idx_bshd = q_indexer.permute(1, 0, 2, 3).contiguous().float() + k_idx_bsd = k_indexer.permute(1, 0, 2).contiguous().float() + w_bsh = weights.permute(1, 0, 2).contiguous().float() + 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. + from megatron.core.transformer.experimental_attention_variant.csa_kernels import ( + _dsa_fwd_flash_mla, + _indexer_topk_bshd, + _kl_loss_from_dense_scores, + _sbhd_to_bshd_indexer_inputs, + ) + + # Run indexer + FlashMLA to capture the same ``lse_indexer`` the fused + # path consumes internally. + effective_topk = min(s['indexer_topk'], s['n_comp']) + q_idx_bshd_bf, k_idx_bsd_bf, _, w_bsh_scaled_bf = _sbhd_to_bshd_indexer_inputs( + q_indexer, k_indexer, weights, s['indexer_softmax_scale'] + ) + topk_indices_cmp, _ = _indexer_topk_bshd( + q_idx_bshd_bf, k_idx_bsd_bf, w_bsh_scaled_bf, effective_topk, s['ratio'] + ) + compress_topk_idxs = torch.where(topk_indices_cmp >= 0, topk_indices_cmp + kv_offset, -1) + combined_local = torch.cat([compress_topk_idxs, win_idxs], dim=-1) + global_idxs = local_to_global_flat(combined_local, s['b'], s['skv']) + q_flat = query.reshape(s['sq'] * s['b'], s['np_'], s['d']) + kv_flat = kv_full.reshape(s['skv'] * s['b'], s['d']) + _, _, lse_indexer = _dsa_fwd_flash_mla( + q_flat, + kv_flat, + global_idxs, + s['softmax_scale'], + attn_sink=attn_sink, + topk_length=None, + indexer_topk=effective_topk, + ) + 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. The fused path + # calls ``_compute_dense_indexer_score`` with ``w_bsh_scaled`` (already + # multiplied by ``indexer_softmax_scale``) AND passes + # ``indexer_softmax_scale`` again as the kernel's ``sm_scale``, + # double-applying the factor (apparent bug in + # ``fused_indexer_sparse_attn`` in ``csa_kernels.py``. Mirror + # that here so the reference matches the fused-path output; revisit + # if the upstream pre-scale + kernel-scale duplication is fixed. + 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'] ** 2 + ) + 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 + ) + assert torch.allclose(indexer_loss, loss_ref, atol=5e-2, rtol=1e-1), ( + f"actual = {indexer_loss.item():.6f}, ref = {loss_ref.item():.6f}, " + f"abs diff = {(indexer_loss - loss_ref).abs().item():.3e}" + ) + + +# =========================================================================== +# THD packed-sequence path +# =========================================================================== + + +def _make_cu_seqlens(seg_lens, device='cpu'): + """Build a ``(B+1,)`` int32 cu_seqlens tensor from a list of segment lengths.""" + return torch.tensor( + [0] + list(torch.tensor(seg_lens, dtype=torch.int64).cumsum(0).tolist()), + dtype=torch.int32, + device=device, + ) + + +class TestThdPureHelpers: + """THD-only pure-Python helpers (no GPU kernels required). + + Covers: + + * ``batch_of_row`` — searchsorted-style ``row → segment`` lookup. + * ``local_to_global_flat`` THD branch — ``cu_seqlens_q/kv`` shift. + * ``build_flat_topk_idxs`` THD branch — ``cu_seqlens_q/kv`` propagation. + * ``_compute_dense_attn_lse`` — THD/segmented-BSHD numerical parity. + """ + + # ---- batch_of_row -------------------------------------------------- + + @pytest.mark.parametrize( + "seg_lens, total_q, expected", + [ + ([3, 3, 3], None, [0, 0, 0, 1, 1, 1, 2, 2, 2]), + ([2, 5, 1, 4], None, [0, 0, 1, 1, 1, 1, 1, 2, 3, 3, 3, 3]), + ([0, 3, 0, 2], None, [1, 1, 1, 3, 3]), + ([5, 5], 7, [0, 0, 0, 0, 0, 1, 1]), + ], + ids=["uniform", "variable", "empty_segment", "total_q_override"], + ) + def test_batch_of_row(self, seg_lens, total_q, expected): + """Row-to-segment mapping for uniform, variable, empty, and truncated cases.""" + cu = _make_cu_seqlens(seg_lens) + bo = batch_of_row(cu, total_q=total_q) if total_q else batch_of_row(cu) + assert bo.tolist() == expected + assert bo.dtype == torch.int64 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def testbatch_of_row_cpu_cuda_parity(self): + """CPU and CUDA executions produce identical results.""" + cu = _make_cu_seqlens([4, 2, 6]) + cpu_out = batch_of_row(cu) + cuda_out = batch_of_row(cu.cuda()) + assert torch.equal(cpu_out, cuda_out.cpu()) + + def test_dense_attn_lse_thd_matches_segmented_bshd(self): + """Packed multi-segment LSE matches independent BSHD segments.""" + torch.manual_seed(123) + q_lens = [5, 7] + kv_lens = [2, 3] + num_heads = 4 + num_kv_heads = 2 + head_dim = 8 + ratio = 2 + qhead_per_kv_head = num_heads // num_kv_heads + + q_thd = torch.randn(sum(q_lens), num_heads, head_dim) + k_thd = torch.randn(sum(kv_lens), num_kv_heads, head_dim) + cu_q = _make_cu_seqlens(q_lens) + cu_kv = _make_cu_seqlens(kv_lens) + actual = dk._compute_dense_attn_lse( + q_thd, + k_thd, + softmax_scale=head_dim**-0.5, + qhead_per_kv_head=qhead_per_kv_head, + ratio=ratio, + cu_seqlens_q=cu_q, + cu_seqlens_kv=cu_kv, + max_seqlen_kv=max(kv_lens), + ) + + expected = [] + q_start = 0 + kv_start = 0 + for q_len, kv_len in zip(q_lens, kv_lens): + expected.append( + dk._compute_dense_attn_lse( + q_thd[q_start : q_start + q_len].unsqueeze(0), + k_thd[kv_start : kv_start + kv_len].unsqueeze(0), + softmax_scale=head_dim**-0.5, + qhead_per_kv_head=qhead_per_kv_head, + ratio=ratio, + ).squeeze(0) + ) + q_start += q_len + kv_start += kv_len + + torch.testing.assert_close(actual, torch.cat(expected)) + + def test_dense_attn_lse_thd_honors_q_causal_offsets(self): + """Local CP rows use their sequence-relative global query positions.""" + torch.manual_seed(321) + q_lens = [2, 3] + kv_lens = [3, 3] + q_causal_offsets = torch.tensor([4, 2], dtype=torch.int32) + num_heads = 4 + num_kv_heads = 2 + head_dim = 8 + ratio = 2 + qhead_per_kv_head = num_heads // num_kv_heads + softmax_scale = head_dim**-0.5 + + q_thd = torch.randn(sum(q_lens), num_heads, head_dim) + k_thd = torch.randn(sum(kv_lens), num_kv_heads, head_dim) + actual = dk._compute_dense_attn_lse( + q_thd, + k_thd, + softmax_scale=softmax_scale, + qhead_per_kv_head=qhead_per_kv_head, + ratio=ratio, + cu_seqlens_q=_make_cu_seqlens(q_lens), + cu_seqlens_kv=_make_cu_seqlens(kv_lens), + max_seqlen_kv=max(kv_lens), + q_causal_offsets=q_causal_offsets, + ) + + expected = [] + q_start = 0 + kv_start = 0 + for segment, (q_len, kv_len) in enumerate(zip(q_lens, kv_lens)): + for q_position in range(q_len): + visible_k = min((int(q_causal_offsets[segment]) + q_position + 1) // ratio, kv_len) + row_lse = [] + for head in range(num_heads): + kv_head = head // qhead_per_kv_head + scores = ( + q_thd[q_start + q_position, head] + * k_thd[kv_start : kv_start + visible_k, kv_head] + ).sum(dim=-1) + row_lse.append(torch.logsumexp(scores * softmax_scale, dim=-1)) + expected.append(torch.stack(row_lse)) + q_start += q_len + kv_start += kv_len + + torch.testing.assert_close(actual, torch.stack(expected)) + + # ---- local_to_global_flat THD branch -------------------------------- + + def test_local_to_global_flat_thd_basic(self): + """THD branch: ``global[i, k] = local[i, k] + cu_seqlens_kv[batch_of_row[i]]``.""" + # Two segments: q lengths [2, 3]; kv lengths [4, 5] (uneven). + cu_q = _make_cu_seqlens([2, 3]) + cu_kv = _make_cu_seqlens([4, 5]) + # local indices: 5 rows × 3 topk; values are per-segment-LOCAL kv ids. + local = torch.tensor( + [ + [0, 1, 2], # seg 0 row 0 → offset 0 + [3, 0, -1], # seg 0 row 1 → offset 0; -1 preserved + [0, 4, 2], # seg 1 row 0 → offset 4 + [1, -1, 3], # seg 1 row 1 → offset 4 + [4, 0, 1], # seg 1 row 2 → offset 4 + ], + dtype=torch.int32, + ) + + out = local_to_global_flat(local, batch_size=-1, cu_seqlens_q=cu_q, cu_seqlens_kv=cu_kv) + expected = torch.tensor( + [[0, 1, 2], [3, 0, -1], [4, 8, 6], [5, -1, 7], [8, 4, 5]], dtype=torch.int32 + ) + assert out.shape == expected.shape + assert out.dtype == torch.int32 + assert torch.equal(out, expected) + + @pytest.mark.parametrize( + "cu_q_segs, cu_kv_segs, match", + [([2, 2], [3, 3, 3], "cu_seqlens"), ([2, 2], None, "must both be provided")], + ids=["shape_mismatch", "xor_cu_seqlens"], + ) + def test_local_to_global_flat_thd_validation(self, cu_q_segs, cu_kv_segs, match): + """Mismatched shapes or supplying only one of cu_seqlens_q/kv raises.""" + local = torch.zeros((4, 2), dtype=torch.int32) + cu_q = _make_cu_seqlens(cu_q_segs) + cu_kv = _make_cu_seqlens(cu_kv_segs) if cu_kv_segs is not None else None + with pytest.raises(ValueError, match=match): + local_to_global_flat(local, -1, cu_seqlens_q=cu_q, cu_seqlens_kv=cu_kv) + + # ---- build_flat_topk_idxs THD branch -------------------------------- + + def test_build_flat_topk_idxs_thd_non_compact(self): + """THD non-compact: concat groups along topk, then THD globalize.""" + cu_q = _make_cu_seqlens([2, 2]) # total_q = 4 + cu_kv = _make_cu_seqlens([3, 3]) # cu_kv = [0, 3, 6] + # Two groups: window-like (2 topk) and compress-like (3 topk). + win = torch.tensor([[0, 1], [1, 2], [0, 1], [-1, 2]], dtype=torch.int32) + cmp_ = torch.tensor([[0, 1, -1], [-1, 0, 1], [0, 2, 1], [1, 0, 2]], dtype=torch.int32) + + flat, length = build_flat_topk_idxs( + win, cmp_, batch_size=-1, cu_seqlens_q=cu_q, cu_seqlens_kv=cu_kv + ) + + # Manual reference: cat then THD globalize (offset = cu_kv[batch_of_row]). + cat = torch.cat([win, cmp_], dim=-1) + expected = local_to_global_flat(cat, -1, cu_seqlens_q=cu_q, cu_seqlens_kv=cu_kv) + assert torch.equal(flat, expected) + assert length is None # non-compact + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_build_flat_topk_idxs_thd_compact_cpu_cuda_parity(self): + """Compact THD path: CPU fallback (CUDA tensors path through the + cuDNN ``compactify`` wrapper if available, else PyTorch fallback) + must produce the same packed valid-first layout in both cases. + """ + # Force CPU fallback by disabling _DSA so compact runs in PyTorch. + saved = dk._DSA + dk._DSA = None + try: + cu_q = _make_cu_seqlens([3, 2]) + cu_kv = _make_cu_seqlens([4, 5]) + local = torch.tensor( + [[0, -1, 2, -1], [1, 2, -1, 0], [-1, 1, -1, 3], [0, 1, 2, -1], [-1, -1, 4, 0]], + dtype=torch.int32, + ) + flat_cpu, len_cpu = build_flat_topk_idxs( + local, batch_size=-1, compact=True, cu_seqlens_q=cu_q, cu_seqlens_kv=cu_kv + ) + flat_cuda, len_cuda = build_flat_topk_idxs( + local.cuda(), + batch_size=-1, + compact=True, + cu_seqlens_q=cu_q.cuda(), + cu_seqlens_kv=cu_kv.cuda(), + ) + assert torch.equal(flat_cpu, flat_cuda.cpu()) + assert torch.equal(len_cpu, len_cuda.cpu()) + # Sanity: valid count per row matches input mask count. + n_valid_per_row = (local >= 0).sum(dim=-1).int() + assert torch.equal(len_cpu, n_valid_per_row) + finally: + dk._DSA = saved + + +class TestThdWrapperDispatchAndValidation: + """THD-mode dispatch + missing-kwarg validation for the three + public layout-aware wrappers: ``indexer_topk``, ``dsa_sparse_attn``, + ``fused_indexer_sparse_attn``. + + All tests are mock-based or shape-only — no real CUDA kernels. + They verify two contracts: + + * **Dispatch**: passing ``cu_seqlens_q`` (or ``is_thd=True``) routes + the wrapper through the THD code path of its underlying kernel + core (as opposed to the SBHD path). + * **Validation**: when the THD path is requested but a required + companion kwarg is missing, the wrapper raises ``ValueError`` + upfront with a clear message (fail-fast, before any kernel + invocation). + + Each section below covers one wrapper. + """ + + # ===================================================================== + # indexer_topk + # ===================================================================== + + def _make_indexer_topk_thd_inputs(self, device='cuda'): + # Two segments, total_q=5, total_k=4 (compressed-K is shorter than Q + # because the indexer K ratio is 4× by default). + cu_q = _make_cu_seqlens([3, 2], device=device) + cu_kv = _make_cu_seqlens([2, 2], device=device) + total_q, total_k = 5, 4 + idx_nh, idx_hd = 4, 64 + q = torch.randn(total_q, idx_nh, idx_hd, dtype=torch.bfloat16, device=device) + k = torch.randn(total_k, idx_hd, dtype=torch.bfloat16, device=device) + w = torch.randn(total_q, idx_nh, dtype=torch.bfloat16, device=device) + return q, k, w, cu_q, cu_kv + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_indexer_topk_thd_dispatch_calls_thd_kernel_path(self, reset_lazy_kernel_state): + """Passing ``cu_seqlens_q`` routes through + ``_DSA.indexer_forward_wrapper`` with ``cu_seqlens_q/k`` + + ``max_seqlen_q/k`` kwargs (THD kernel mode), as opposed to the + positional-only BSHD call. + """ + 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') + + def fake_indexer_forward(q_thd, k_thd, w_thd, ratio, **kwargs): + # Verify THD kwargs were forwarded. + assert 'cu_seqlens_q' in kwargs and kwargs['cu_seqlens_q'] is cu_q + 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 + fake_dsa.indexer_top_k_wrapper.side_effect = lambda scores_flat, seq_lens, **kw: { + 'indices': torch.zeros( + scores_flat.shape[0], kw['top_k'], dtype=torch.int32, device=scores_flat.device + ) + } + dk._DSA = fake_dsa + + topk_idxs, topk_len = indexer_topk( + q, + k, + w, + topk=2, + ratio=4, + indexer_softmax_scale=128**-0.5, + cu_seqlens_q=cu_q, + 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) + # ===================================================================== + + @pytest.mark.parametrize( + "query_shape, kv_shape, match", + [ + ((4, 2, 4, 64), (8, 64), "THD dsa_sparse_attn expects query"), + ((4, 4, 64), (8, 2, 64), "THD dsa_sparse_attn expects kv"), + ], + ids=["query_wrong_ndim", "kv_wrong_ndim"], + ) + def test_dsa_sparse_attn_thd_wrong_ndim_raises(self, query_shape, kv_shape, match): + """THD mode requires ``query.ndim == 3`` and ``kv.ndim == 2``.""" + query = torch.zeros(*query_shape, dtype=torch.bfloat16) + kv = torch.zeros(*kv_shape, dtype=torch.bfloat16) + attn_sink = torch.zeros(4, dtype=torch.float32) + topk = torch.zeros(query_shape[0], 2, dtype=torch.int32) + with pytest.raises(ValueError, match=match): + 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): + """THD inputs flow into ``SparseAttnFunc`` unchanged (no SBHD + reshape) and the output is ``(total_q, np * d_v)``. + """ + total_q, np_, d, d_v = 6, 4, 64, 512 + n_kv = 8 + query = torch.randn(total_q, np_, d, dtype=torch.bfloat16, device='cuda') + kv = torch.randn(n_kv, d, dtype=torch.bfloat16, device='cuda') + attn_sink = torch.zeros(np_, dtype=torch.float32, device='cuda') + topk = torch.zeros(total_q, 2, dtype=torch.int32, device='cuda') + + flash_stub = _make_flash_mla_stub(d_v=d_v) + dk._flash_mla_sparse_fwd = flash_stub + + out = dsa_sparse_attn(query, kv, attn_sink, topk, softmax_scale=0.125, is_thd=True) + assert out.shape == (total_q, np_ * d_v) + # The FlashMLA stub was called with the unmodified flat tensors. + # _dsa_fwd_flash_mla passes (q, kv_3d, indices, softmax_scale) positionally; + # kv is unsqueezed to (n_kv, 1, d) before reaching the kernel. + call_q, call_kv_3d = flash_stub.call_args.args[0], flash_stub.call_args.args[1] + assert call_q.shape == (total_q, np_, d) + assert call_kv_3d.squeeze(1).shape == (n_kv, d) + + # ===================================================================== + # fused_indexer_sparse_attn (validation only — real-kernel parity is + # covered by TestRealKernelFusedIndexerSparseAttnThd below) + # ===================================================================== + + def _fused_common_thd_kwargs(self): + return dict( + cu_seqlens_q=_make_cu_seqlens([2, 2]), + cu_seqlens_kv=_make_cu_seqlens([2, 2]), + cu_seqlens_kv_full=_make_cu_seqlens([3, 3]), + cu_seqlens_compressed_idx=_make_cu_seqlens([1, 1]), + max_seqlen_q=2, + max_seqlen_compressed_idx=1, + ) + + def _fused_dummy_thd_inputs(self): + total_q, np_, d = 4, 4, 64 + total_kv_full = 6 + total_comp_idx = 2 + idx_nh, idx_hd = 4, 64 + return dict( + query=torch.zeros(total_q, np_, d, dtype=torch.bfloat16), + kv_full=torch.zeros(total_kv_full, d, dtype=torch.bfloat16), + attn_sink=torch.zeros(np_, dtype=torch.float32), + window_idxs=torch.zeros(total_q, 2, dtype=torch.int32), + q_indexer=torch.zeros(total_q, idx_nh, idx_hd, dtype=torch.bfloat16), + k_indexer=torch.zeros(total_comp_idx, idx_hd, dtype=torch.bfloat16), + weights=torch.zeros(total_q, idx_nh, dtype=torch.bfloat16), + ) + + @pytest.mark.parametrize( + "missing", + [ + 'cu_seqlens_kv', + 'cu_seqlens_kv_full', + 'cu_seqlens_compressed_idx', + 'max_seqlen_q', + 'max_seqlen_compressed_idx', + ], + ) + def test_fused_indexer_sparse_attn_thd_missing_kwarg_raises(self, missing): + """All five THD-companion kwargs are required when ``cu_seqlens_q`` + is supplied; a missing one raises ``ValueError`` upfront. + + (No ``test_thd_sparse_loss_raises``: sparse-loss is supported in + THD via flat-global topk ids — see + ``TestRealKernelFusedIndexerSparseAttnThd``.) + """ + kwargs = self._fused_common_thd_kwargs() + kwargs[missing] = None + inputs = self._fused_dummy_thd_inputs() + with pytest.raises(ValueError, match="THD mode requires"): + fused_indexer_sparse_attn( + **inputs, indexer_topk=2, ratio=4, softmax_scale=0.125, **kwargs + ) + + +# --------------------------------------------------------------------------- +# Real-kernel THD parity +# --------------------------------------------------------------------------- + + +class TestRealKernelFusedIndexerSparseAttnThd: + """End-to-end parity for Path B in THD mode (both loss variants): + real cuDNN score-recompute + indexer-backward kernels + real FlashMLA, + compared to the equivalent SBHD invocation on the same data with B=1. + + For a single-segment THD batch (``cu_seqlens_q = [0, sq]``) the THD + pipeline produces a numerically equivalent loss to the SBHD pipeline + with ``b=1`` on the same tensors — both go through the same + underlying cuDNN kernels, differing only in the layout-glue around + them. The sparse-loss THD path additionally exercises + :func:`local_to_global_flat` (over ``cu_seqlens_compressed_idx``) + and the ``topk_indices_global=True`` flag wiring. + """ + + SHAPES = dict( + sq=128, + np_=64, + d=512, + skv=640, + n_comp=512, + # cudnn DSA (dense_)indexer_backward kernels require heads >= 64. + idx_nh=64, + idx_hd=128, + indexer_topk=512, + ratio=4, + win_topk=8, + softmax_scale=512**-0.5, + indexer_softmax_scale=128**-0.5, + ) + + @pytest.mark.parametrize('sparse_loss', [False, True], ids=['dense_loss', 'sparse_loss']) + def test_thd_single_segment_matches_sbhd_b1(self, sparse_loss, reset_lazy_kernel_state): + """B=1 THD invocation should match the equivalent SBHD-b=1 call + 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) + s = self.SHAPES + torch.manual_seed(0) + dev = 'cuda' + loss_coeff = 0.5 + b = 1 + + # Common inputs (SBHD layout — single batch). + query_sbhd = torch.randn(s['sq'], b, s['np_'], s['d'], dtype=torch.bfloat16, device=dev) + kv_full_sbhd = torch.randn(s['skv'], b, s['d'], dtype=torch.bfloat16, device=dev) + attn_sink = torch.zeros(s['np_'], dtype=torch.float32, device=dev) + torch.manual_seed(1) + win_idxs_sbhd = torch.randint( + 0, s['sq'], (b, s['sq'], s['win_topk']), dtype=torch.int32, device=dev + ) + q_indexer_sbhd = torch.randn( + s['sq'], b, s['idx_nh'], s['idx_hd'], dtype=torch.bfloat16, device=dev + ) + k_indexer_sbhd = torch.randn(s['n_comp'], b, s['idx_hd'], dtype=torch.bfloat16, device=dev) + weights_sbhd = torch.randn(s['sq'], b, s['idx_nh'], dtype=torch.bfloat16, device=dev) + kv_offset = s['skv'] - s['n_comp'] + + # ---- SBHD reference -------------------------------------------------- + _, loss_sbhd = fused_indexer_sparse_attn( + query_sbhd, + kv_full_sbhd, + attn_sink, + win_idxs_sbhd, + q_indexer_sbhd, + k_indexer_sbhd, + weights_sbhd, + indexer_topk=s['indexer_topk'], + ratio=s['ratio'], + softmax_scale=s['softmax_scale'], + indexer_softmax_scale=s['indexer_softmax_scale'], + loss_coeff=loss_coeff, + sparse_loss=sparse_loss, + kv_offset=kv_offset, + ) + + # ---- THD equivalent -------------------------------------------------- + # Reshape: SBHD (sq, 1, ...) -> THD flat (sq, ...). + # kv_full SBHD layout is [kv (sq), compressed (n_comp)] in dim 0; + # the THD analogue is [kv (sq), compressed (n_comp)] per-segment. + query_thd = query_sbhd.squeeze(1) # (sq, np, d) + kv_full_thd = kv_full_sbhd.squeeze(1) # (skv, d) + win_idxs_thd = win_idxs_sbhd.squeeze(0) # (sq, win_topk) + q_indexer_thd = q_indexer_sbhd.squeeze(1) # (sq, idx_nh, idx_hd) + k_indexer_thd = k_indexer_sbhd.squeeze(1) # (n_comp, idx_hd) + weights_thd = weights_sbhd.squeeze(1) # (sq, idx_nh) + + # Single-segment cu_seqlens (B=1): total_q == sq. + cu_q = _make_cu_seqlens([s['sq']], device=dev) + cu_kv = _make_cu_seqlens([kv_offset], device=dev) + cu_kv_full = _make_cu_seqlens([s['skv']], device=dev) + # Indexer K is per-segment compressed-only (n_comp positions). + cu_comp_idx = _make_cu_seqlens([s['n_comp']], device=dev) + + # B=1: per-segment [kv, compressed] layout collapses to a single + # contiguous slice — same kv_offset as the SBHD case. + compressed_kv_thd = kv_full_thd[kv_offset:] + _, loss_thd = fused_indexer_sparse_attn( + query_thd, + kv_full_thd, + attn_sink, + win_idxs_thd, + q_indexer_thd, + k_indexer_thd, + weights_thd, + indexer_topk=s['indexer_topk'], + ratio=s['ratio'], + softmax_scale=s['softmax_scale'], + indexer_softmax_scale=s['indexer_softmax_scale'], + loss_coeff=loss_coeff, + sparse_loss=sparse_loss, + kv_offset=0, # ignored in THD + cu_seqlens_q=cu_q, + cu_seqlens_kv=cu_kv, + cu_seqlens_kv_full=cu_kv_full, + cu_seqlens_compressed_idx=cu_comp_idx, + max_seqlen_q=s['sq'], + max_seqlen_compressed_idx=s['n_comp'], + compressed_kv=compressed_kv_thd, + ) + + # SBHD and THD share the same underlying kernels; for B=1 the + # numerical paths are identical up to topk-ordering ties in the + # indexer's radix top-K, which can shift a few scores at the + # boundary. Use the same tolerance as the SBHD-vs-PyTorch test. + assert torch.allclose(loss_thd, loss_sbhd, atol=5e-2, rtol=1e-1), ( + f"sparse_loss={sparse_loss}: thd = {loss_thd.item():.6f}, " + f"sbhd = {loss_sbhd.item():.6f}, " + f"abs diff = {(loss_thd - loss_sbhd).abs().item():.3e}" + ) + + +# --------------------------------------------------------------------------- +# THD padding-row masking: cu_seqlens_q_unpadded excludes padding from loss +# --------------------------------------------------------------------------- + + +class TestThdPaddingRowMasking: + """Verify that **per-segment** padding rows do NOT contribute to the + indexer KL loss when ``cu_seqlens_q_unpadded`` is supplied. + """ + + SHAPES = dict( + np_=64, + d=512, + idx_nh=64, + idx_hd=128, + indexer_topk=512, + ratio=4, + win_topk=8, + softmax_scale=512**-0.5, + indexer_softmax_scale=128**-0.5, + ) + # 3 sequences with per-segment padding. + SEG_LENS_REAL = [60, 44, 720] # real token counts per sequence + SEG_LENS_PADDED = [64, 48, 1024] # padded to multiple of 4 + + @staticmethod + def _build_multi_seg_inputs(seg_lens, shapes, dev, *, seed=42): + """Build THD multi-segment inputs for fused_indexer_sparse_attn. + + Each segment has its own original KV (len = seg_len) and compressed + KV (len = seg_len // ratio), concatenated per-segment in kv_full. + """ + torch.manual_seed(seed) + s = shapes + ratio = s['ratio'] + total_q = sum(seg_lens) + comp_lens = [sl // ratio for sl in seg_lens] + total_comp = sum(comp_lens) + kv_full_seg_lens = [sl + cl for sl, cl in zip(seg_lens, comp_lens)] + total_kv_full = sum(kv_full_seg_lens) + max_seqlen_q = max(seg_lens) + max_comp = max(comp_lens) if comp_lens else 0 + + query = torch.randn(total_q, s['np_'], s['d'], dtype=torch.bfloat16, device=dev) + kv_full = torch.randn(total_kv_full, s['d'], dtype=torch.bfloat16, device=dev) + attn_sink = torch.zeros(s['np_'], dtype=torch.float32, device=dev) + + # Per-segment local window indices. + win_idxs = torch.zeros(total_q, s['win_topk'], dtype=torch.int32, device=dev) + offset = 0 + for sl in seg_lens: + if sl > 0: + win_idxs[offset : offset + sl] = torch.randint( + 0, sl, (sl, s['win_topk']), dtype=torch.int32, device=dev + ) + offset += sl + + q_indexer = torch.randn(total_q, s['idx_nh'], s['idx_hd'], dtype=torch.bfloat16, device=dev) + k_indexer = torch.randn(total_comp, s['idx_hd'], dtype=torch.bfloat16, device=dev) + weights = torch.randn(total_q, s['idx_nh'], dtype=torch.bfloat16, device=dev) + + compressed_parts = [] + kv_offset = 0 + for sl, cl in zip(seg_lens, comp_lens): + compressed_parts.append(kv_full[kv_offset + sl : kv_offset + sl + cl]) + kv_offset += sl + cl + compressed_kv = torch.cat(compressed_parts, dim=0) if compressed_parts else kv_full[:0] + + cu_q = _make_cu_seqlens(seg_lens, device=dev) + cu_kv = _make_cu_seqlens(seg_lens, device=dev) + cu_kv_full = _make_cu_seqlens(kv_full_seg_lens, device=dev) + cu_comp = _make_cu_seqlens(comp_lens, device=dev) + + return dict( + query=query, + kv_full=kv_full, + attn_sink=attn_sink, + win_idxs=win_idxs, + q_indexer=q_indexer, + k_indexer=k_indexer, + weights=weights, + compressed_kv=compressed_kv, + cu_q=cu_q, + cu_kv=cu_kv, + cu_kv_full=cu_kv_full, + cu_comp=cu_comp, + total_q=total_q, + max_seqlen_q=max_seqlen_q, + max_comp=max_comp, + seg_lens=seg_lens, + comp_lens=comp_lens, + ) + + @staticmethod + def _build_per_seg_padded(real_inputs, padded_seg_lens, shapes, dev, *, fill_pad_random=False): + """Expand real inputs to a per-segment-padded layout. + + Each segment is expanded from its real length to its padded length + (padding rows inserted at the tail of each segment). + + Returns (padded_inputs_dict, cu_q_unpadded). + """ + s = shapes + ratio = s['ratio'] + r = real_inputs + real_seg_lens = r['seg_lens'] + num_segs = len(real_seg_lens) + total_q_padded = sum(padded_seg_lens) + comp_lens_padded = [pl // ratio for pl in padded_seg_lens] + total_comp_padded = sum(comp_lens_padded) + kv_full_seg_lens_padded = [pl + cl for pl, cl in zip(padded_seg_lens, comp_lens_padded)] + total_kv_full_padded = sum(kv_full_seg_lens_padded) + + fill_fn = torch.randn if fill_pad_random else torch.zeros + + # Build padded Q-side tensors by scattering real data into padded slots. + query_pad = fill_fn(total_q_padded, s['np_'], s['d'], dtype=torch.bfloat16, device=dev) + q_idx_pad = fill_fn( + total_q_padded, s['idx_nh'], s['idx_hd'], dtype=torch.bfloat16, device=dev + ) + w_pad = fill_fn(total_q_padded, s['idx_nh'], dtype=torch.bfloat16, device=dev) + win_pad = torch.zeros(total_q_padded, s['win_topk'], dtype=torch.int32, device=dev) + + real_offset = 0 + pad_offset = 0 + for i in range(num_segs): + rl = real_seg_lens[i] + pl = padded_seg_lens[i] + query_pad[pad_offset : pad_offset + rl] = r['query'][real_offset : real_offset + rl] + q_idx_pad[pad_offset : pad_offset + rl] = r['q_indexer'][real_offset : real_offset + rl] + w_pad[pad_offset : pad_offset + rl] = r['weights'][real_offset : real_offset + rl] + win_pad[pad_offset : pad_offset + rl] = r['win_idxs'][real_offset : real_offset + rl] + real_offset += rl + pad_offset += pl + + # Build padded K-side (compressed indexer K). + k_idx_pad = fill_fn(total_comp_padded, s['idx_hd'], dtype=torch.bfloat16, device=dev) + comp_kv_pad = fill_fn(total_comp_padded, s['d'], dtype=torch.bfloat16, device=dev) + real_comp_offset = 0 + pad_comp_offset = 0 + for i in range(num_segs): + rcl = r['comp_lens'][i] + pcl = comp_lens_padded[i] + k_idx_pad[pad_comp_offset : pad_comp_offset + rcl] = r['k_indexer'][ + real_comp_offset : real_comp_offset + rcl + ] + comp_kv_pad[pad_comp_offset : pad_comp_offset + rcl] = r['compressed_kv'][ + real_comp_offset : real_comp_offset + rcl + ] + real_comp_offset += rcl + pad_comp_offset += pcl + + # Build padded kv_full: per-segment [orig_kv (padded_len), compressed (padded_comp)]. + kv_full_pad = fill_fn(total_kv_full_padded, s['d'], dtype=torch.bfloat16, device=dev) + real_kv_offset = 0 + pad_kv_offset = 0 + real_comp_offset2 = 0 + pad_comp_offset2 = 0 + for i in range(num_segs): + rl = real_seg_lens[i] + pl = padded_seg_lens[i] + rcl = r['comp_lens'][i] + pcl = comp_lens_padded[i] + # Copy real orig-KV rows. + src_start = sum(s + c for s, c in zip(real_seg_lens[:i], r['comp_lens'][:i])) + kv_full_pad[pad_kv_offset : pad_kv_offset + rl] = r['kv_full'][ + src_start : src_start + rl + ] + # Copy real compressed rows. + kv_full_pad[pad_kv_offset + pl : pad_kv_offset + pl + rcl] = r['kv_full'][ + src_start + rl : src_start + rl + rcl + ] + pad_kv_offset += pl + pcl + + cu_q_padded = _make_cu_seqlens(padded_seg_lens, device=dev) + cu_kv_padded = _make_cu_seqlens(padded_seg_lens, device=dev) + cu_kv_full_padded = _make_cu_seqlens(kv_full_seg_lens_padded, device=dev) + cu_comp_padded = _make_cu_seqlens(comp_lens_padded, device=dev) + + # Unpadded cu_seqlens: cumulative REAL lengths within the padded layout. + cu_q_unpadded = _make_cu_seqlens(list(real_seg_lens), device=dev) + + max_seqlen_q_padded = max(padded_seg_lens) + max_comp_padded = max(comp_lens_padded) + + return ( + dict( + query=query_pad, + kv_full=kv_full_pad, + attn_sink=r['attn_sink'], + win_idxs=win_pad, + q_indexer=q_idx_pad, + k_indexer=k_idx_pad, + weights=w_pad, + compressed_kv=comp_kv_pad, + cu_q=cu_q_padded, + cu_kv=cu_kv_padded, + cu_kv_full=cu_kv_full_padded, + cu_comp=cu_comp_padded, + total_q=total_q_padded, + max_seqlen_q=max_seqlen_q_padded, + max_comp=max_comp_padded, + ), + cu_q_unpadded, + ) + + def _run_fused( + self, inputs, shapes, *, sparse_loss, loss_coeff=0.5, cu_seqlens_q_unpadded=None + ): + """Run fused_indexer_sparse_attn with the given inputs dict.""" + s = shapes + i = inputs + return fused_indexer_sparse_attn( + i['query'], + i['kv_full'], + i['attn_sink'], + i['win_idxs'], + i['q_indexer'], + i['k_indexer'], + i['weights'], + indexer_topk=s['indexer_topk'], + ratio=s['ratio'], + softmax_scale=s['softmax_scale'], + indexer_softmax_scale=s['indexer_softmax_scale'], + loss_coeff=loss_coeff, + sparse_loss=sparse_loss, + kv_offset=0, + cu_seqlens_q=i['cu_q'], + cu_seqlens_kv=i['cu_kv'], + cu_seqlens_kv_full=i['cu_kv_full'], + cu_seqlens_compressed_idx=i['cu_comp'], + max_seqlen_q=i['max_seqlen_q'], + max_seqlen_compressed_idx=i['max_comp'], + compressed_kv=i['compressed_kv'], + cu_seqlens_q_unpadded=cu_seqlens_q_unpadded, + # Per-token (sum) reduction — the real training path. Padding + # rows contribute 0 to the sum, so the loss is padding-invariant + # by construction; the global token divisor is applied later by + # DSAIndexerLossAutoScaler.set_loss_scale. Mean reduction would + # instead divide by the padded row count and dilute the loss. + calculate_per_token_loss=True, + ) + + @pytest.mark.parametrize('sparse_loss', [False, True], ids=['dense_loss', 'sparse_loss']) + def test_per_seg_padding_excluded_from_loss(self, sparse_loss, reset_lazy_kernel_state): + """Per-segment padding rows should not contribute to indexer KL. + + Strategy: compute loss on tightly-packed real data (no padding), + 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) + dev = 'cuda' + + # Baseline: tightly packed (real lengths only, no padding). + real = self._build_multi_seg_inputs(self.SEG_LENS_REAL, self.SHAPES, dev) + _, loss_no_pad = self._run_fused(real, self.SHAPES, sparse_loss=sparse_loss) + + # Padded: each segment expanded to padded length (zeros in padding). + padded, cu_q_unpadded = self._build_per_seg_padded( + real, self.SEG_LENS_PADDED, self.SHAPES, dev, fill_pad_random=False + ) + _, loss_with_pad = self._run_fused( + padded, self.SHAPES, sparse_loss=sparse_loss, cu_seqlens_q_unpadded=cu_q_unpadded + ) + + assert torch.allclose(loss_with_pad, loss_no_pad, atol=5e-2, rtol=1e-1), ( + f"sparse_loss={sparse_loss}: padded = {loss_with_pad.item():.6f}, " + f"no_pad = {loss_no_pad.item():.6f}, " + f"abs diff = {(loss_with_pad - loss_no_pad).abs().item():.3e}" + ) + + @pytest.mark.parametrize('sparse_loss', [False, True], ids=['dense_loss', 'sparse_loss']) + def test_per_seg_padding_unmasked_corrupts_loss(self, sparse_loss, reset_lazy_kernel_state): + """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) + dev = 'cuda' + + real = self._build_multi_seg_inputs(self.SEG_LENS_REAL, self.SHAPES, dev) + _, loss_no_pad = self._run_fused(real, self.SHAPES, sparse_loss=sparse_loss) + + # Padded with RANDOM noise in per-segment padding slots. + padded, _ = self._build_per_seg_padded( + real, self.SEG_LENS_PADDED, self.SHAPES, dev, fill_pad_random=True + ) + _, loss_unmasked = self._run_fused(padded, self.SHAPES, sparse_loss=sparse_loss) + + assert not torch.allclose(loss_unmasked, loss_no_pad, atol=5e-2, rtol=1e-1), ( + f"sparse_loss={sparse_loss}: unmasked loss ({loss_unmasked.item():.6f}) should " + f"differ from no-pad loss ({loss_no_pad.item():.6f}) since per-segment " + "padding has random data producing non-zero KL" + ) + + @pytest.mark.parametrize('sparse_loss', [False, True], ids=['dense_loss', 'sparse_loss']) + def test_per_seg_padding_grads_are_zeroed(self, sparse_loss, reset_lazy_kernel_state): + """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) + dev = 'cuda' + + # ---- Unpadded baseline (reference grads) ----------------------------- + real = self._build_multi_seg_inputs(self.SEG_LENS_REAL, self.SHAPES, dev) + real['q_indexer'] = real['q_indexer'].detach().requires_grad_(True) + real['k_indexer'] = real['k_indexer'].detach().requires_grad_(True) + real['weights'] = real['weights'].detach().requires_grad_(True) + + _, loss_real = self._run_fused(real, self.SHAPES, sparse_loss=sparse_loss) + loss_real.backward() + grad_q_real = real['q_indexer'].grad.detach().clone() + grad_w_real = real['weights'].grad.detach().clone() + + # ---- Padded run with masking ----------------------------------------- + real_nograd = self._build_multi_seg_inputs(self.SEG_LENS_REAL, self.SHAPES, dev) + padded, cu_q_unpadded = self._build_per_seg_padded( + real_nograd, self.SEG_LENS_PADDED, self.SHAPES, dev, fill_pad_random=True + ) + + padded['q_indexer'] = padded['q_indexer'].detach().requires_grad_(True) + padded['k_indexer'] = padded['k_indexer'].detach().requires_grad_(True) + padded['weights'] = padded['weights'].detach().requires_grad_(True) + + _, indexer_loss = self._run_fused( + padded, self.SHAPES, sparse_loss=sparse_loss, cu_seqlens_q_unpadded=cu_q_unpadded + ) + indexer_loss.backward() + + # ---- Identify real and padding positions ----------------------------- + real_positions = [] + pad_positions = [] + pad_offset = 0 + for rl, pl in zip(self.SEG_LENS_REAL, self.SEG_LENS_PADDED): + for pos in range(rl): + real_positions.append(pad_offset + pos) + for pos in range(rl, pl): + pad_positions.append(pad_offset + pos) + pad_offset += pl + real_positions = torch.tensor(real_positions, dtype=torch.long, device=dev) + pad_positions = torch.tensor(pad_positions, dtype=torch.long, device=dev) + + # ---- Assert: padding positions have zero grad ------------------------ + pad_grad_q = padded['q_indexer'].grad[pad_positions] + assert torch.all(pad_grad_q == 0), ( + f"q_indexer grad at per-segment padding positions should be zero, " + f"got max abs = {pad_grad_q.abs().max().item():.3e}" + ) + pad_grad_w = padded['weights'].grad[pad_positions] + assert torch.all(pad_grad_w == 0), ( + f"weights grad at per-segment padding positions should be zero, " + f"got max abs = {pad_grad_w.abs().max().item():.3e}" + ) + + # ---- Assert: real positions match unpadded baseline grads ------------- + # The cuDNN indexer-backward kernel is non-deterministic (a config + # compared against itself shows per-element grad diffs ~= the max grad + # magnitude), so an element-wise allclose is unachievable. Instead + # compare *direction* via a global (flattened) cosine similarity: + # padded-vs-baseline measures ~0.997 while the same-config noise floor + # is ~0.9995, so >0.99 robustly confirms the masking preserves the + # real-token gradients while still catching a genuinely corrupted mask. + # Per-row cosine is unusable here: causal-masked early rows have + # all-zero grads (cosine vs a zero vector is 0). + def _grad_cos_sim(a, b): + return torch.nn.functional.cosine_similarity( + a.flatten().float(), b.flatten().float(), dim=0 + ) + + cos_q = _grad_cos_sim(padded['q_indexer'].grad[real_positions], grad_q_real) + assert cos_q > 0.99, ( + f"q_indexer grad at real positions should align with unpadded " + f"baseline, cosine similarity = {cos_q.item():.6f}" + ) + cos_w = _grad_cos_sim(padded['weights'].grad[real_positions], grad_w_real) + assert cos_w > 0.99, ( + f"weights grad at real positions should align with unpadded " + f"baseline, cosine similarity = {cos_w.item():.6f}" + ) + + +class TestFusedIndexerSparseAttnFromTopk: + """Focused contract tests for the CP caller-supplied-top-k autograd path.""" + + @pytest.mark.parametrize("sparse_loss", [False, True], ids=["dense", "sparse"]) + def test_offsets_padding_mask_and_backward(self, sparse_loss, monkeypatch): + total_q, num_heads, head_dim = 3, 2, 4 + total_kv, total_comp = 6, 4 + index_heads, index_dim, topk = 2, 3, 2 + + query = torch.randn(total_q, num_heads, head_dim, requires_grad=True) + kv_full = torch.randn(total_kv, head_dim, requires_grad=True) + attn_sink = torch.randn(num_heads, requires_grad=True) + topk_idxs = torch.tensor([[0, 1], [2, 3], [4, 5]], dtype=torch.int32) + q_indexer = torch.randn(total_q, index_heads, index_dim, requires_grad=True) + k_indexer = torch.randn(total_comp, index_dim, requires_grad=True) + weights = torch.randn(total_q, index_heads, requires_grad=True) + indexer_topk_idxs = torch.tensor([[0, 1], [1, 2], [2, 3]], dtype=torch.int32) + compressed_kv = torch.randn(total_comp, head_dim) + cu_seqlens_q = torch.tensor([0, total_q], dtype=torch.int32) + cu_seqlens_k = torch.tensor([0, total_comp], dtype=torch.int32) + q_causal_offsets = torch.tensor([1], dtype=torch.int32) + q_padding_mask = torch.tensor([False, True, False]) + + def fake_flash(q, kv, caller_topk, _softmax_scale, **kwargs): + assert torch.equal(caller_topk, topk_idxs) + assert kwargs["indexer_topk"] == topk + return ( + torch.zeros_like(q), + torch.zeros(total_q, num_heads), + torch.zeros(total_q, num_heads), + ) + + monkeypatch.setattr(dk, "_dsa_fwd_flash_mla", fake_flash) + fake_dsa = MagicMock() + monkeypatch.setattr(dk, "_DSA", fake_dsa) + + def fake_sparse_attn_backward(q, kv, _out, _grad, _lse, sink, caller_topk, **kwargs): + assert torch.equal(caller_topk, topk_idxs) + return {"dq": torch.ones_like(q), "dkv": torch.ones_like(kv), "d_sink": torch.ones_like(sink)} + + fake_dsa.sparse_attention_backward_wrapper.side_effect = fake_sparse_attn_backward + + if sparse_loss: + + def fake_sparse_score(_q, _k, _w, caller_indexer_topk, **kwargs): + assert torch.equal( + caller_indexer_topk.squeeze(0)[q_padding_mask], + torch.full((1, topk), -1, dtype=torch.int32), + ) + return {"predict": torch.full((1, total_q, topk), 1.0 / topk)} + + def fake_target(_q, _k, _lse, caller_indexer_topk, *_args, **_kwargs): + assert torch.equal( + caller_indexer_topk[q_padding_mask], + torch.full((1, topk), -1, dtype=torch.int32), + ) + return torch.full((total_q, topk), 1.0 / topk) + + fake_dsa.sparse_indexer_score_recompute_wrapper.side_effect = fake_sparse_score + fake_dsa.indexer_backward_wrapper.return_value = { + "d_index_q": torch.full((1, total_q, index_heads, index_dim), 2.0), + "d_index_k": torch.full((1, total_comp, index_dim), 3.0), + "d_weights": torch.full((1, total_q, index_heads), 4.0), + } + monkeypatch.setattr(dk, "_compute_attn_target", fake_target) + else: + + def assert_offsets(kwargs): + assert kwargs["q_causal_offsets"] is q_causal_offsets + + def fake_dense_indexer(*_args, **kwargs): + assert_offsets(kwargs) + return torch.zeros(total_q, total_comp), torch.full( + (total_q,), math.log(total_comp) + ) + + def fake_dense_lse(*_args, **kwargs): + assert_offsets(kwargs) + return torch.zeros(total_q, num_heads) + + def fake_dense_attn_score(*_args, **kwargs): + assert_offsets(kwargs) + return torch.ones(total_q, total_comp), torch.full( + (total_q,), float(total_comp) + ) + + def fake_dense_backward( + _q, _w, _k, attn_score, _attn_l1norm, index_score, index_lse, **kwargs + ): + assert_offsets(kwargs) + assert torch.equal(attn_score[q_padding_mask], torch.zeros(1, total_comp)) + assert torch.equal(index_score[q_padding_mask], torch.zeros(1, total_comp)) + assert torch.equal(index_lse[q_padding_mask], torch.zeros(1)) + return { + "d_index_q": torch.full_like(q_indexer, 2.0), + "d_index_k": torch.full_like(k_indexer, 3.0), + "d_weights": torch.full_like(weights, 4.0), + } + + monkeypatch.setattr(dk, "_compute_dense_indexer_score", fake_dense_indexer) + monkeypatch.setattr(dk, "_compute_dense_attn_lse", fake_dense_lse) + monkeypatch.setattr(dk, "_compute_dense_attn_score", fake_dense_attn_score) + fake_dsa.dense_indexer_backward_wrapper.side_effect = fake_dense_backward + + output, indexer_loss = FusedIndexerSparseAttnFromTopkFunc.apply( + query, + kv_full, + attn_sink, + topk_idxs, + q_indexer, + k_indexer, + weights, + indexer_topk_idxs, + compressed_kv, + 0.5, + 0.25, + 1.0, + float(total_q), + sparse_loss, + 1, + total_comp, + (cu_seqlens_q, cu_seqlens_k, q_causal_offsets), + q_padding_mask, + ) + (output.sum() + indexer_loss).backward() + + assert output.shape == (total_q, num_heads * head_dim) + torch.testing.assert_close(query.grad, torch.ones_like(query)) + torch.testing.assert_close(kv_full.grad, torch.ones_like(kv_full)) + torch.testing.assert_close(attn_sink.grad, torch.ones_like(attn_sink)) + torch.testing.assert_close(q_indexer.grad[~q_padding_mask], torch.full_like(q_indexer.grad[~q_padding_mask], 2.0)) + torch.testing.assert_close(q_indexer.grad[q_padding_mask], torch.zeros_like(q_indexer.grad[q_padding_mask])) + torch.testing.assert_close(k_indexer.grad, torch.full_like(k_indexer, 3.0)) + torch.testing.assert_close(weights.grad[~q_padding_mask], torch.full_like(weights.grad[~q_padding_mask], 4.0)) + torch.testing.assert_close(weights.grad[q_padding_mask], torch.zeros_like(weights.grad[q_padding_mask])) + + +# --------------------------------------------------------------------------- +# Public surface +# --------------------------------------------------------------------------- + + +class TestPublicApi: + """The ``__all__`` list documents the public surface; verify that every + advertised symbol is importable, that the public free functions are + callable, and that the autograd Functions inherit from the right base. + """ + + def test_public_surface(self): + from megatron.core.transformer.experimental_attention_variant import csa_kernels + + for name in csa_kernels.__all__: + assert hasattr(csa_kernels, name), f"__all__ lists {name!r} but it is missing" + + for fn in ( + build_flat_topk_idxs, + local_to_global_flat, + dsa_sparse_attn, + indexer_topk, + fused_indexer_sparse_attn, + ): + assert callable(fn) + + assert issubclass(SparseAttnFunc, torch.autograd.Function) + assert issubclass(FusedIndexerSparseAttnFunc, torch.autograd.Function) + assert issubclass(FusedIndexerSparseAttnFromTopkFunc, torch.autograd.Function) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_backend_tp_sp_parity.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_backend_tp_sp_parity.py index 8d63a9bee11..b6c4e4a1be5 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_backend_tp_sp_parity.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_backend_tp_sp_parity.py @@ -681,3 +681,112 @@ def test_packed_cp_tp2_sequence_parallel_shared_skip_backend_matches_unfused_ref finally: DSAIndexerLossLoggingHelper.clean_loss_in_tracker() Utils.destroy_model_parallel() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_indexer_loss_tracker_grows_for_hybrid_mtp_layer_numbers(): + """Hybrid MTP layers can have a layer number beyond the nominal layer count.""" + DSAIndexerLossLoggingHelper.tracker = {} + try: + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=torch.tensor(2.0, device="cuda"), + layer_number=7, + num_layers=5, + ) + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=torch.tensor(3.0, device="cuda"), + layer_number=9, + num_layers=5, + ) + + values = DSAIndexerLossLoggingHelper.tracker["values"] + assert values.shape == (9,) + torch.testing.assert_close(values[6], torch.tensor(2.0, device="cuda")) + torch.testing.assert_close(values[8], torch.tensor(3.0, device="cuda")) + finally: + DSAIndexerLossLoggingHelper.tracker = {} + + +def test_clean_indexer_loss_tracker_preserves_group_identity_when_requested(): + original_tracker = DSAIndexerLossLoggingHelper.tracker + reduce_group, avg_group = object(), object() + DSAIndexerLossLoggingHelper.tracker = { + "values": torch.ones(2), + "reduce_group": reduce_group, + "avg_group": avg_group, + } + try: + DSAIndexerLossLoggingHelper.clean_loss_in_tracker(preserve_groups=True) + + tracker = DSAIndexerLossLoggingHelper.tracker + assert torch.count_nonzero(tracker["values"]) == 0 + assert tracker["reduce_group"] is reduce_group + assert tracker["avg_group"] is avg_group + + DSAIndexerLossLoggingHelper.clean_loss_in_tracker() + assert tracker["reduce_group"] is None + assert tracker["avg_group"] is None + finally: + DSAIndexerLossLoggingHelper.tracker = original_tracker + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_indexer_metrics_average_only_ratio4_layers(monkeypatch: pytest.MonkeyPatch): + """Window and compressed-only layers must not dilute the indexer-loss average.""" + DSAIndexerLossLoggingHelper.tracker = { + "values": torch.tensor([0.0, 3.0, 0.0, 0.0, 0.0], device="cuda") + } + monkeypatch.setattr(DSAIndexerLossLoggingHelper, "reduce_loss_in_tracker", lambda **_: None) + total_loss_dict = {} + try: + DSAIndexerLossLoggingHelper.track_indexer_metrics( + loss_scale=0.5, + iteration=1, + writer=None, + pg_collection=object(), + total_loss_dict=total_loss_dict, + num_layers=5, + csa_compress_ratios=[0, 4, 128, 0, 0], + ) + + torch.testing.assert_close( + total_loss_dict["indexer loss"], torch.tensor(1.5, device="cuda") + ) + finally: + DSAIndexerLossLoggingHelper.tracker = {} + + +def test_indexer_metrics_reduce_across_pipeline_rank_without_indexer(): + """Every pipeline rank must join indexer loss reduction, even without a local indexer.""" + if Utils.world_size < 2: + pytest.skip("Cross-pipeline indexer reduction requires at least two distributed ranks") + + Utils.initialize_model_parallel(tensor_model_parallel_size=1, pipeline_model_parallel_size=2) + DSAIndexerLossLoggingHelper.tracker = {} + try: + num_layers = 5 + if parallel_state.get_pipeline_model_parallel_rank() == 0: + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=torch.tensor(3.0, device="cuda"), layer_number=2, num_layers=num_layers + ) + + total_loss_dict = {} + DSAIndexerLossLoggingHelper.track_indexer_metrics( + loss_scale=0.5, + iteration=1, + writer=None, + pg_collection=ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['pp', 'dp'] + ), + total_loss_dict=total_loss_dict, + num_layers=num_layers, + csa_compress_ratios=[0, 4, 128, 0, 0], + ) + + torch.testing.assert_close( + total_loss_dict["indexer loss"], torch.tensor(1.5, device="cuda") + ) + assert torch.count_nonzero(DSAIndexerLossLoggingHelper.tracker["values"]) == 0 + finally: + DSAIndexerLossLoggingHelper.tracker = {} + Utils.destroy_model_parallel() diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py new file mode 100644 index 00000000000..e980bfab5d4 --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py @@ -0,0 +1,840 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from megatron.core.extensions.transformer_engine import HAVE_TE +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.transformer_config import MLATransformerConfig +from tests.unit_tests.test_utilities import Utils + +try: + from fast_hadamard_transform import hadamard_transform as _hadamard_transform + + HAVE_HADAMARD = True +except ImportError: + HAVE_HADAMARD = False + _hadamard_transform = None + +_SEED = 42 + + +def _mock_hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: + return x * scale + + +@pytest.fixture(autouse=True) +def patch_hadamard_if_needed(): + """Patch hadamard_transform in dsa/csa modules if the library is not installed.""" + if not HAVE_HADAMARD: + with ( + patch( + 'megatron.core.transformer.experimental_attention_variant.dsa.hadamard_transform', + _mock_hadamard_transform, + ), + patch( + 'megatron.core.transformer.experimental_attention_variant.csa.rotate_activation', + lambda x: x * (x.size(-1) ** -0.5), + ), + ): + yield + else: + yield + + +# --------------------------------------------------------------------------- +# Config / spec helpers +# --------------------------------------------------------------------------- + + +def _make_config( + num_layers=4, + hidden_size=256, + num_attention_heads=16, + v_head_dim=64, + qk_pos_emb_head_dim=32, + q_lora_rank=64, + o_groups=8, + o_lora_rank=64, + csa_compress_ratios=None, + csa_window_size=8, + tensor_model_parallel_size=1, + sequence_parallel=False, + dsa_indexer_n_heads=8, + dsa_indexer_head_dim=64, + dsa_indexer_topk=8, + dsa_indexer_loss_coeff=0.0, + **extra_config_kwargs, +): + """Create an MLATransformerConfig for DSv4 hybrid attention tests.""" + if csa_compress_ratios is None: + csa_compress_ratios = [0, 4, 128, 4] + return MLATransformerConfig( + num_layers=num_layers, + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + use_cpu_initialization=True, + bf16=True, + params_dtype=torch.bfloat16, + add_bias_linear=False, + tensor_model_parallel_size=tensor_model_parallel_size, + sequence_parallel=sequence_parallel, + q_lora_rank=q_lora_rank, + kv_lora_rank=v_head_dim - qk_pos_emb_head_dim, + qk_head_dim=v_head_dim - qk_pos_emb_head_dim, + qk_pos_emb_head_dim=qk_pos_emb_head_dim, + v_head_dim=v_head_dim, + o_groups=o_groups, + o_lora_rank=o_lora_rank, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + multi_latent_attention=True, + experimental_attention_variant='dsv4_hybrid', + csa_compress_ratios=csa_compress_ratios, + csa_window_size=csa_window_size, + dsa_indexer_n_heads=dsa_indexer_n_heads, + dsa_indexer_head_dim=dsa_indexer_head_dim, + dsa_indexer_topk=dsa_indexer_topk, + dsa_indexer_loss_coeff=dsa_indexer_loss_coeff, + **extra_config_kwargs, + ) + + +def _make_attention_spec(config): + """Build the full DSv4HybridSelfAttention ModuleSpec using the canonical spec builder.""" + from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider + from megatron.core.transformer.experimental_attention_variant.dsv4_module_specs import ( + get_dsv4_hybrid_module_spec_for_backend, + ) + + return get_dsv4_hybrid_module_spec_for_backend(config=config, backend=TESpecProvider()) + + +def test_module_spec_is_built_from_explicit_backend(): + """The neutral spec builder should use only its explicitly supplied backend.""" + from megatron.core.transformer.experimental_attention_variant.csa import ( + CompressedSparseAttention, + Compressor, + CSAIndexer, + ) + from megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention import ( + DSv4HybridSelfAttention, + ) + from megatron.core.transformer.experimental_attention_variant.dsv4_module_specs import ( + get_dsv4_hybrid_module_spec_for_backend, + ) + + class Linear: + pass + + class ColumnParallelLinear: + pass + + class RowParallelLinear: + pass + + class Norm: + pass + + class Backend: + def linear(self): + return Linear + + def column_parallel_linear(self): + return ColumnParallelLinear + + def row_parallel_linear(self): + return RowParallelLinear + + def layer_norm(self, rms_norm=False, for_qk=False, has_residual=False): + return Norm + + spec = get_dsv4_hybrid_module_spec_for_backend(_make_config(), Backend()) + + assert spec.module is DSv4HybridSelfAttention + assert spec.submodules.linear_q_down_proj is Linear + assert spec.submodules.linear_q_up_proj is ColumnParallelLinear + assert spec.submodules.linear_kv_proj is ColumnParallelLinear + assert spec.submodules.linear_proj is RowParallelLinear + assert spec.submodules.core_attention.module is CompressedSparseAttention + assert spec.submodules.core_attention.submodules.compressor.module is Compressor + assert spec.submodules.core_attention.submodules.indexer.module is CSAIndexer + + +def test_config_includes_mtp_ratio_and_derives_dimensions(): + """DSv4 config should account for MTP and derive its shared Q/KV content width.""" + config = _make_config(num_layers=2, mtp_num_layers=1, csa_compress_ratios=[0, 4, 128]) + + expected_content_dim = config.v_head_dim - config.qk_pos_emb_head_dim + assert config.qk_head_dim == expected_content_dim + assert config.kv_lora_rank == expected_content_dim + assert config.hetereogenous_dist_checkpoint is True + + +def test_config_accepts_multi_layer_hybrid_mtp_ratio_list(): + config = _make_config(num_layers=2, mtp_num_layers=1, csa_compress_ratios=[0, 4, 128, 0]) + assert config.csa_compress_ratios == [0, 4, 128, 0] + + +def test_hybrid_dsv4_stack_spec_assigns_symbol_ratios(): + from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_dsv4_stack_spec + + stack_spec = hybrid_dsv4_stack_spec(_make_config()) + submodules = stack_spec.submodules + + assert "compress_ratio" not in submodules.dsa_layer.submodules.self_attention.params + assert submodules.csa_layer.submodules.self_attention.params["compress_ratio"] == 4 + assert submodules.hca_layer.submodules.self_attention.params["compress_ratio"] == 128 + assert submodules.window_layer.submodules.self_attention.params["compress_ratio"] == 0 + + +def test_config_accepts_static_context_parallelism(): + with patch( + "megatron.core.transformer.transformer_config.is_te_min_version", return_value=True + ): + config = _make_config( + context_parallel_size=2, + sequence_packing_scheduler="dp_balanced", + cp_partition_mode="contiguous", + ) + + assert config.context_parallel_size == 2 + assert config.cp_partition_mode == "contiguous" + + +@pytest.mark.parametrize( + ("config_kwargs", "match"), + [ + ({"context_parallel_size": 2}, "sequence_packing_scheduler"), + ( + {"context_parallel_size": 2, "sequence_packing_scheduler": "dp_balanced"}, + "cp_partition_mode='contiguous'", + ), + ({"hybrid_context_parallel": True}, "dynamic per-microbatch context parallelism"), + ({"cp_partition_mode": "interleaved"}, "Unsupported cp_partition_mode"), + ], + ids=["missing_scheduler", "zigzag", "dynamic_cp", "invalid_partition"], +) +def test_config_rejects_unsupported_context_parallelism(config_kwargs, match): + with pytest.raises(ValueError, match=match): + _make_config(**config_kwargs) + + +@pytest.mark.parametrize( + ("packed_kwargs", "match"), + [ + ( + {"qkv_format": "thd", "local_cp_size": 2, "cp_partition_mode": "contiguous"}, + "per-microbatch context-parallel groups", + ), + ({"qkv_format": "sbhd", "cp_partition_mode": "contiguous"}, "qkv_format='thd'"), + ({"qkv_format": "thd", "cp_partition_mode": "zigzag"}, "contiguous CP partition"), + ], + ids=["local_cp", "non_thd", "zigzag"], +) +def test_forward_rejects_unsupported_context_parallel_inputs(packed_kwargs, match): + from megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention import ( + DSv4HybridAttention, + ) + + cp_group = SimpleNamespace(size=lambda: 2) + attention = SimpleNamespace(pg_collection=SimpleNamespace(cp=cp_group)) + + with pytest.raises(ValueError, match=match): + DSv4HybridAttention.forward( + attention, + hidden_states=None, + attention_mask=None, + packed_seq_params=PackedSeqParams(**packed_kwargs), + ) + + +def _build_attention(config, layer_number, pg_collection, **kwargs): + """Instantiate a DSv4HybridSelfAttention from config.""" + from megatron.core.transformer.spec_utils import build_module + + spec = _make_attention_spec(config) + return build_module( + spec, config=config, layer_number=layer_number, pg_collection=pg_collection, **kwargs + ) + + +# =========================================================================== +# Constructor tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridAttentionConstructor: + """Test construction of DSv4HybridSelfAttention in the supported TP=1 configuration.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + yield + Utils.destroy_model_parallel() + + def test_basic_construction(self): + """Verify the layer builds and has the expected sub-modules.""" + from megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention import ( + DSv4HybridSelfAttention, + ) + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + config = _make_config() + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention(config, layer_number=1, pg_collection=pg) + + assert isinstance(attn, DSv4HybridSelfAttention) + assert hasattr(attn, 'linear_q_down_proj') + assert hasattr(attn, 'linear_q_up_proj') + assert hasattr(attn, 'linear_kv_proj') + assert hasattr(attn, 'linear_proj') + assert hasattr(attn, 'linear_o_group_proj') + assert hasattr(attn, 'core_attention') + assert hasattr(attn, 'q_layernorm') + assert hasattr(attn, 'kv_layernorm') + + def test_q_head_dim_equals_v_head_dim(self): + """q_head_dim must equal v_head_dim for DSv4 hybrid.""" + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + config = _make_config() + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention(config, layer_number=1, pg_collection=pg) + + assert attn.q_head_dim == config.v_head_dim + + def test_current_main_constructor_kwargs(self): + """Current TransformerLayer forwards module names and pipeline offsets.""" + config = _make_config() + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention( + config, + layer_number=1, + pg_collection=pg, + pp_layer_offset=0, + name="decoder.layers.0.self_attention", + ) + + assert attn._pp_layer_offset == 0 + + @pytest.mark.parametrize("layer_number", [1, 2, 3, 4]) + def test_rope_base_varies_with_compress_ratio(self, layer_number): + """Layers with compress_ratio > 1 should use csa_compress_rotary_base.""" + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + ratios = [0, 4, 128, 4] + config = _make_config(csa_compress_ratios=ratios) + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention(config, layer_number=layer_number, pg_collection=pg) + + ratio = ratios[layer_number - 1] + if ratio > 1: + expected_base = config.csa_compress_rotary_base + else: + expected_base = config.rotary_base + + # inv_freq is derived from rotary_base; verify the correct base was used + dim = config.qk_pos_emb_head_dim + recomputed_inv_freq = 1.0 / ( + expected_base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim) + ) + assert torch.allclose( + attn.rotary_pos_emb.inv_freq.cpu(), recomputed_inv_freq, rtol=1e-5, atol=1e-5 + ) + + +# =========================================================================== +# Forward / backward tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridAttentionForwardBackward: + """Test forward and backward passes of DSv4HybridSelfAttention.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + cls = request.cls + cls.config = _make_config(dsa_indexer_loss_coeff=1.0) + cls.pg = ProcessGroupCollection.use_mpu_process_groups() + + yield + Utils.destroy_model_parallel() + + @pytest.mark.parametrize("layer_number", [1, 2, 3, 4]) + def test_forward_output_shape(self, layer_number): + """Forward should produce [sq, b, hidden_size] output.""" + seq_len = 256 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention( + self.config, layer_number=layer_number, pg_collection=self.pg + ).cuda() + + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + output, bias = attn(hidden_states=hidden, attention_mask=None) + + assert output.shape == (seq_len, batch_size, self.config.hidden_size) + assert output.dtype == torch.bfloat16 + assert not torch.isnan(output).any() + + @pytest.mark.parametrize("layer_number", [1, 2]) + def test_backward_gradient_flow(self, layer_number): + """Backward should produce gradients for all trainable parameters.""" + seq_len = 256 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention( + self.config, layer_number=layer_number, pg_collection=self.pg + ).cuda() + attn.train() + + hidden = ( + torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) + + output, bias = attn(hidden_states=hidden, attention_mask=None) + loss = output.sum() + loss.backward() + + assert hidden.grad is not None, "No gradient on hidden_states" + for name, param in attn.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"No gradient for parameter {name}" + + def test_eval_mode(self): + """Forward should work in eval mode.""" + seq_len = 128 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention(self.config, layer_number=1, pg_collection=self.pg).cuda() + attn.eval() + + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + with torch.no_grad(): + output, bias = attn(hidden_states=hidden, attention_mask=None) + + assert output.shape == (seq_len, batch_size, self.config.hidden_size) + assert not torch.isnan(output).any() + + def test_different_seq_lengths(self): + """Forward should handle various sequence lengths.""" + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention(self.config, layer_number=2, pg_collection=self.pg).cuda() + + for seq_len in [64, 128, 256]: + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + output, bias = attn(hidden_states=hidden, attention_mask=None) + assert output.shape == (seq_len, batch_size, self.config.hidden_size) + + +# =========================================================================== +# get_query_key_value_tensors tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridQKV: + """Test get_query_key_value_tensors internals.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + cls = request.cls + cls.config = _make_config() + cls.pg = ProcessGroupCollection.use_mpu_process_groups() + + yield + Utils.destroy_model_parallel() + + def test_qkv_shapes(self): + """Query, key, value should have correct shapes.""" + seq_len = 64 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention(self.config, layer_number=1, pg_collection=self.pg).cuda() + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + q, k, v, q_compressed, kv_compressed = attn.get_query_key_value_tensors(hidden) + + n_heads = self.config.num_attention_heads + v_dim = self.config.v_head_dim + + assert q.shape == (seq_len, batch_size, n_heads, v_dim) + # key and value are single-head (MQA-style) with an extra head dim + assert k.shape[-1] == v_dim + assert v.shape[-1] == v_dim + + def test_key_equals_value(self): + """In the wkv path, key and value should be the same tensor.""" + seq_len = 64 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention(self.config, layer_number=1, pg_collection=self.pg).cuda() + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + q, k, v, _, _ = attn.get_query_key_value_tensors(hidden) + assert torch.equal(k, v), "key and value should be identical in wkv path" + + +# =========================================================================== +# Grouped output projection tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridGroupedOutput: + """Test that grouped output projection (wo_a) parameters are created.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + yield + Utils.destroy_model_parallel() + + def test_o_group_proj_shape(self): + """linear_o_group_proj should have the correct shape.""" + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + o_groups = 8 + o_lora_rank = 64 + config = _make_config(o_groups=o_groups, o_lora_rank=o_lora_rank) + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention(config, layer_number=1, pg_collection=pg) + + expected_out = o_groups * o_lora_rank + expected_in = (config.v_head_dim * config.num_attention_heads) // o_groups + assert attn.linear_o_group_proj.shape == (expected_out, expected_in) + assert attn.linear_o_group_proj.requires_grad + + +# =========================================================================== +# apply_rope_fusion tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridRopeFusion: + """Test that apply_rope_fusion=True works for both yarn and non-yarn layers. + + DSv4 Hybrid uses YarnRotaryEmbedding for layers with compress_ratio > 1 + and standard RotaryEmbedding for layers with compress_ratio <= 1. The + fused RoPE path must obtain cos/sin from both embedding classes via + get_cached_cos_sin. + + compress_ratios=[0, 4, 128, 4]: layer 1 has ratio 0 (standard + RotaryEmbedding), layers 2-4 have ratio > 1 (YarnRotaryEmbedding). + """ + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + cls = request.cls + cls.pg = ProcessGroupCollection.use_mpu_process_groups() + + yield + Utils.destroy_model_parallel() + + def test_rope_fusion_forward_backward_parity(self): + """Fused RoPE forward/backward succeeds and matches the unfused path.""" + seq_len = 128 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + fused_config = _make_config(apply_rope_fusion=True) + attn_fused = _build_attention(fused_config, layer_number=4, pg_collection=self.pg).cuda() + attn_fused.train() + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + unfused_config = _make_config(apply_rope_fusion=False) + attn_unfused = _build_attention( + unfused_config, layer_number=4, pg_collection=self.pg + ).cuda() + attn_unfused.train() + + hidden = torch.randn( + seq_len, batch_size, fused_config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + out_fused, _ = attn_fused(hidden_states=hidden, attention_mask=None) + out_unfused, _ = attn_unfused(hidden_states=hidden, attention_mask=None) + + assert out_fused.shape == (seq_len, batch_size, fused_config.hidden_size) + assert torch.isfinite(out_fused).all() + # The remaining difference is bf16 accumulation order between the fused + # Triton kernel and eager PyTorch operations. + torch.testing.assert_close(out_fused, out_unfused, atol=3e-2, rtol=3e-2) + + hidden_fused = hidden.detach().clone().requires_grad_(True) + hidden_unfused = hidden.detach().clone().requires_grad_(True) + + attn_fused(hidden_states=hidden_fused, attention_mask=None)[0].sum().backward() + attn_unfused(hidden_states=hidden_unfused, attention_mask=None)[0].sum().backward() + + assert hidden_fused.grad is not None + for name, param in attn_fused.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"No gradient for parameter {name}" + +# =========================================================================== +# THD packed-sequence end-to-end +# =========================================================================== +# +# Closes the highest-level integration gap: even though the CSA THD path +# is independently tested in test_attention_variant_csa.py +# (TestCompressedSparseAttentionThd), the DSv4HybridSelfAttention module +# adds its own THD-aware glue around CSA — packed_seq_params propagation +# through get_query_key_value_tensors, the output reshape at line ~290 +# (``core_attn_out.reshape(total, 1, -1)`` to recover the 3-D contract), +# and the inverse-RoPE call with cu_seqlens. These tests verify the full +# DSv4Hybrid forward/backward works end-to-end for each ``compress_ratio``. + + + +def _make_thd_packed_seq_params(seg_lens, device='cuda'): + """Build ``PackedSeqParams(qkv_format='thd', ...)`` for self-attention + (``cu_seqlens_q == cu_seqlens_kv``) from a list of per-segment lengths. + """ + cu_seqlens = torch.tensor( + [0] + list(torch.tensor(seg_lens, dtype=torch.int64).cumsum(0).tolist()), + dtype=torch.int32, + device=device, + ) + max_len = int(max(seg_lens)) if seg_lens else 0 + return PackedSeqParams( + cu_seqlens_q=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_kv_padded=cu_seqlens, + max_seqlen_q=max_len, + max_seqlen_kv=max_len, + qkv_format='thd', + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridAttentionThd: + """End-to-end THD forward/backward of :class:`DSv4HybridSelfAttention` + across all configured ``compress_ratio`` values (0/4/128). + + Each test runs a multi-segment THD batch through the full + ``attn(hidden_states, packed_seq_params=...)`` pipeline and verifies: + + * Output shape ``(total_tokens, 1, hidden_size)`` — the layer + re-adds the dummy ``b=1`` axis at line ~293 of + ``deepseek_v4_hybrid_attention.py``. + * No NaN. + * (Backward test) grads flow on ``hidden_states`` and every + learnable parameter. + + This is the highest-level integration test for THD; lower-level + parity vs SBHD is established by ``TestCompressedSparseAttentionThd`` + in ``test_attention_variant_csa.py``. + """ + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + cls = request.cls + # ``csa_compress_ratios=[0, 4, 128, 4]`` → layer_number ∈ {1,2,3,4} + # cover all three CSA ratios (0 = window-only; 4 = full + # indexer+compressed; 128 = compressor-only, no indexer). + cls.config = _make_config(dsa_indexer_loss_coeff=1.0) + cls.pg = ProcessGroupCollection.use_mpu_process_groups() + + yield + Utils.destroy_model_parallel() + + @pytest.mark.parametrize( + "layer_number", + [1, 2, 3, 4], + ids=[ + "ratio_0_window_only", # layer 1 → ratio=0 + "ratio_4_with_indexer", # layer 2 → ratio=4 + "ratio_128_compressor_only", # layer 3 → ratio=128 + "ratio_4_with_indexer_alt", # layer 4 → ratio=4 + ], + ) + def test_thd_forward_output_shape(self, layer_number): + """THD forward through DSv4HybridSelfAttention produces + ``(total_tokens, 1, hidden_size)`` output, no NaN, for every + configured ``compress_ratio``. + """ + seg_lens = [128, 96, 64] # multi-segment, varied lengths + total = sum(seg_lens) + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention( + self.config, layer_number=layer_number, pg_collection=self.pg + ).cuda() + attn.eval() + + # THD hidden_states shape is ``(total_tokens, 1, hidden_size)`` + # per the DSv4 hybrid contract. + hidden = torch.randn(total, 1, self.config.hidden_size, dtype=torch.bfloat16, device='cuda') + packed = _make_thd_packed_seq_params(seg_lens) + + with torch.no_grad(): + output, _bias = attn( + hidden_states=hidden, attention_mask=None, packed_seq_params=packed + ) + + assert output.shape == (total, 1, self.config.hidden_size), ( + f"layer {layer_number}: shape {tuple(output.shape)} != " + f"expected {(total, 1, self.config.hidden_size)}" + ) + assert output.dtype == torch.bfloat16 + assert not torch.isnan(output).any(), f"layer {layer_number}: NaN in THD forward output" + + @pytest.mark.parametrize( + "layer_number", + [1, 2], # ratio=0 (window-only) and ratio=4 (full indexer pipeline) + ids=["ratio_0_window_only", "ratio_4_with_indexer"], + ) + def test_thd_backward_gradient_flow(self, layer_number): + """THD backward produces grads on ``hidden_states`` and every + learnable parameter (covers the full indexer-loss path in Path + B THD when ``layer_number=2`` triggers ratio=4). + """ + seg_lens = [128, 96] + total = sum(seg_lens) + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention( + self.config, layer_number=layer_number, pg_collection=self.pg + ).cuda() + attn.train() + + hidden = torch.randn( + total, 1, self.config.hidden_size, dtype=torch.bfloat16, device='cuda' + ).requires_grad_(True) + packed = _make_thd_packed_seq_params(seg_lens) + + output, _bias = attn(hidden_states=hidden, attention_mask=None, packed_seq_params=packed) + output.sum().backward() + + assert hidden.grad is not None, "no grad on hidden_states" + assert not torch.isnan(hidden.grad).any() + for name, param in attn.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"no grad on {name}" + assert not torch.isnan(param.grad).any(), f"NaN grad on {name}" + + def test_thd_single_segment_matches_sbhd_b1(self): + """B=1 single-segment THD output matches the SBHD-b=1 output on + identical hidden states (sanity check that the DSv4Hybrid + THD-vs-SBHD glue doesn't silently change the math). + + Uses ``layer_number=1`` (ratio=0, window-only) for determinism — + no indexer top-K tie-breaking nondeterminism. + """ + layer_number = 1 # ratio=0 → window-only path, no cuDNN topk + sq = 128 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention( + self.config, layer_number=layer_number, pg_collection=self.pg + ).cuda() + attn.eval() + + hidden = torch.randn(sq, 1, self.config.hidden_size, dtype=torch.bfloat16, device='cuda') + + with torch.no_grad(): + out_sbhd, _ = attn(hidden_states=hidden, attention_mask=None, packed_seq_params=None) + packed = _make_thd_packed_seq_params([sq]) + out_thd, _ = attn(hidden_states=hidden, attention_mask=None, packed_seq_params=packed) + + assert out_sbhd.shape == out_thd.shape + # Generous tol: the full DSv4Hybrid forward chains many bf16 ops + # (QKV down/up proj, RoPE, attn, output proj). We're testing for + # no plumbing bug, not bit-exactness. + assert torch.allclose(out_sbhd.float(), out_thd.float(), atol=5e-2, rtol=5e-2), ( + f"B=1 SBHD/THD parity failed: max abs diff = " + f"{(out_sbhd.float() - out_thd.float()).abs().max().item():.4e}" + ) 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..95be09b655a --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py @@ -0,0 +1,835 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import gc +import os +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 +_DSV4_CP_TEST_VARIANT = "flash" + +# 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 + +# 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 _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() + 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() + + @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, monkeypatch + ): + """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) + + project_weights_calls = 0 + if layer_number == 2: + original_project_weights = cp_attn.core_attention.indexer._project_weights + + def tracked_project_weights(x): + nonlocal project_weights_calls + project_weights_calls += 1 + return original_project_weights(x) + + monkeypatch.setattr( + cp_attn.core_attention.indexer, "_project_weights", tracked_project_weights + ) + + 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 + ) + if layer_number == 2: + assert project_weights_calls == 1 + 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_cp_mla_up_proj_recompute_matches_eager(self, apply_rope_fusion): + """Selective MLA recompute must retain the configured 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) + + 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.pg).cuda() + recompute_attn = _build_attention( + recompute_config, layer_number=2, pg_collection=self.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"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() 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 new file mode 100644 index 00000000000..f9f7b74a77f --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_native_parity.py @@ -0,0 +1,1408 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import gc +import math + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from megatron.core.extensions.transformer_engine import HAVE_TE +from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.experimental_attention_variant.dsa import DSAIndexerLossAutoScaler +from megatron.core.transformer.experimental_attention_variant.dsv4_module_specs import ( + get_dsv4_hybrid_module_spec_for_backend, +) +from megatron.core.transformer.spec_utils import build_module +from megatron.core.transformer.transformer_config import MLATransformerConfig +from megatron.core.utils import init_method_normal, scaled_init_method_normal +from tests.unit_tests.test_utilities import Utils + +_SEED = 1234 +# Parity tolerances (cosine / tensor-sim drift = 1 - sim), split on two axes: +# +# * fused vs unfused — the fused path exercises the cudnn DSA kernels + Triton +# fused MLA RoPE, whose bf16 numerics (and non-deterministic atomic +# reductions) drift ~an order of magnitude more than the pytorch-eager +# unfused path. ``apply_rope_fusion`` is coupled to ``apply_dsa_kernel_fusion``. +# * forward (the layer ``out``) vs backward (``hidden_grad`` + every param +# grad) — gradients accumulate kernel noise and need looser floors than the +# forward output. +# +# Each constant covers the worst case across the whole parametrization for its +# (path, direction) bucket. Values sit ~1.3-2.5x above the measured worst-case +# drift over the full matrix (variant x ratio x seqlen x segment-layout); the +# forward buckets have wide headroom (the layer output is a well-averaged +# quantity), the backward buckets are near their physical floor: +# * fused-fwd worst ~8e-4 (layer ``out``) -> 2e-3 +# * fused-bwd worst ~1.6e-2 (``core_attention.attn_sink`` — a per-head +# scalar grad with no spatial averaging; the binding param) -> 2e-2 +# * unfused-fwd worst ~7e-4 (layer ``out``) -> 1.5e-3 +# * unfused-bwd worst ~2e-3 (compressor / indexer grads, ratio > 1) -> 3e-3 +# A real positioning/aggregation regression collapses cosine far below these +# floors, so the budgets still trip on genuine bugs. +_FUSED_FWD_SIMILARITY_EPS = 2e-3 +_FUSED_BWD_SIMILARITY_EPS = 2e-2 +_UNFUSED_FWD_SIMILARITY_EPS = 1.5e-3 +_UNFUSED_BWD_SIMILARITY_EPS = 3e-3 + + +@torch.compile +def _native_q_rms_norm(query: torch.Tensor, eps: float) -> torch.Tensor: + return query * torch.rsqrt(query.square().mean(-1, keepdim=True) + eps) + + +_DSV4_VARIANTS = { + "flash": { + "hidden_size": 4096, + "num_attention_heads": 64, + "q_lora_rank": 1024, + "v_head_dim": 512, + "qk_pos_emb_head_dim": 64, + "o_groups": 8, + "o_lora_rank": 1024, + "csa_compress_rotary_base": 40000, + "dsa_indexer_topk": 512, + }, +} + +_DSA_BACKENDS = [ + pytest.param("fused", True, id="fused"), + pytest.param("unfused", False, id="unfused"), +] + + +def _make_config( + variant: str, + compress_ratio: int, + apply_dsa_kernel_fusion: bool = False, + calculate_per_token_loss: bool = False, + dsa_indexer_use_sparse_loss: bool = False, +) -> MLATransformerConfig: + shape = _DSV4_VARIANTS[variant] + mcore_ratio = 0 if compress_ratio == 1 else compress_ratio + qk_head_dim = shape["v_head_dim"] - shape["qk_pos_emb_head_dim"] + config = MLATransformerConfig( + multi_latent_attention=True, + experimental_attention_variant="dsv4_hybrid", + num_layers=1, + hidden_size=shape["hidden_size"], + num_attention_heads=shape["num_attention_heads"], + q_lora_rank=shape["q_lora_rank"], + kv_lora_rank=qk_head_dim, + qk_head_dim=qk_head_dim, + qk_pos_emb_head_dim=shape["qk_pos_emb_head_dim"], + v_head_dim=shape["v_head_dim"], + o_groups=shape["o_groups"], + o_lora_rank=shape["o_lora_rank"], + csa_compress_ratios=[mcore_ratio], + csa_window_size=128, + csa_dense_mode=False, + dsa_indexer_n_heads=64, + dsa_indexer_head_dim=128, + dsa_indexer_topk=shape["dsa_indexer_topk"], + dsa_indexer_loss_coeff=0.01, + dsa_indexer_use_sparse_loss=dsa_indexer_use_sparse_loss, + calculate_per_token_loss=calculate_per_token_loss, + add_bias_linear=False, + bf16=True, + params_dtype=torch.bfloat16, + layernorm_epsilon=1e-6, + normalization="RMSNorm", + qk_layernorm=True, + layernorm_zero_centered_gamma=False, + expert_model_parallel_size=1, + tensor_model_parallel_size=1, + sequence_parallel=False, + context_parallel_size=1, + rope_type="yarn" if apply_dsa_kernel_fusion else "rope", + rotary_base=10000, + rotary_percent=1.0, + csa_compress_rotary_base=shape["csa_compress_rotary_base"], + recompute_granularity=None, + recompute_modules=[], + fine_grained_activation_offloading=False, + gradient_accumulation_fusion=False, + fp8=False, + fp4=False, + init_method=init_method_normal(0.02), + output_layer_init_method=scaled_init_method_normal(0.02, 1, multiplier=2.0), + kv_channels=shape["v_head_dim"], + num_query_groups=shape["num_attention_heads"], + batch_invariant_mode=False, + cache_mla_latents=False, + use_cpu_initialization=True, + perform_initialization=True, + symmetric_ar_type=None, + disable_parameter_transpose_cache=False, + init_model_with_meta_device=False, + delay_wgrad_compute=False, + tp_comm_overlap=False, + softmax_scale=None, + apply_dsa_kernel_fusion=apply_dsa_kernel_fusion, + apply_rope_fusion=apply_dsa_kernel_fusion, + ) + return config + + +def _precompute_freqs_cis( + dim: int, + seqlen: int, + device, + base: float, + *, + original_seq_len: int = 0, + factor: float = 1.0, + beta_fast: float = 32.0, + beta_slow: float = 1.0, +) -> torch.Tensor: + """Precompute the [seq, 1, 1, dim] freqs table used by ``_apply_rotary_emb``. + + Matches the golden DSv4 reference (``Megatron-LM/model.py:precompute_freqs_cis``) + and ``YarnRotaryEmbedding`` semantics: + + * ``original_seq_len > 0`` enables YaRN frequency interpolation between + the ``beta_fast`` / ``beta_slow`` correction-range bounds. Frequencies + below the low boundary are divided by ``factor`` (interpolation); above + the high boundary, freqs pass through (extrapolation); a smooth linear + ramp blends the two in between. + * ``original_seq_len == 0`` reverts to plain RoPE with no scaling — the + window-only branch on the production side. + """ + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)) + if original_seq_len > 0: + + def _correction_dim(num_rotations): + return (dim * math.log(original_seq_len / (num_rotations * 2 * math.pi))) / ( + 2 * math.log(base) + ) + + low = max(int(math.floor(_correction_dim(beta_fast))), 0) + high = min(int(math.ceil(_correction_dim(beta_slow))), dim - 1) + if low == high: + high += 1 # avoid div-by-zero in the ramp + ramp = (torch.arange(dim // 2, dtype=torch.float32, device=device) - low) / (high - low) + smooth = 1.0 - torch.clamp(ramp, 0.0, 1.0) + freqs = freqs / factor * (1.0 - smooth) + freqs * smooth + + t = torch.arange(seqlen, device=device) + freqs = torch.outer(t, freqs) + return torch.cat((freqs, freqs), dim=-1)[:, None, None, :] + + +def _apply_rotary_emb( + x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False +) -> torch.Tensor: + if x.numel() == 0: + return x + freqs = freqs_cis.to(x.device) + if freqs.dim() == x.dim() + 1 and freqs.size(-2) == 1: + freqs = freqs.squeeze(-2) + + rot_dim = freqs.size(-1) + x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:] + x1 = x_rot[..., 0::2] + x2 = x_rot[..., 1::2] + x_rot = torch.cat((x1, x2), dim=-1) + + cos = torch.cos(freqs).to(x_rot.dtype) + sin = torch.sin(freqs).to(x_rot.dtype) + if inverse: + sin = -sin + + rot_half_1, rot_half_2 = torch.chunk(x_rot, 2, dim=-1) + x_rotated = torch.cat((-rot_half_2, rot_half_1), dim=-1) + out = (x_rot * cos) + (x_rotated * sin) + + x1, x2 = torch.chunk(out, 2, dim=-1) + out = torch.stack((x1, x2), dim=-1).flatten(start_dim=-2) + return torch.cat((out, x_pass), dim=-1) + + +def _native_hadamard_transform(x: torch.Tensor) -> torch.Tensor: + n = x.size(-1) + if n <= 0 or n & (n - 1): + raise ValueError(f"Hadamard transform requires power-of-two last dim, got {n}") + dtype = x.dtype + y = x.float() + shape = y.shape + h = 1 + while h < n: + y = y.reshape(*shape[:-1], -1, 2, h) + a = y[..., 0, :] + b = y[..., 1, :] + y = torch.cat((a + b, a - b), dim=-1) + h *= 2 + return (y.reshape(shape) * (n**-0.5)).to(dtype) + + +def _get_window_topk_idxs( + window_size: int, batch_size: int, seqlen: int, device: torch.device +) -> torch.Tensor: + base = torch.arange(seqlen, device=device).unsqueeze(1) + offsets = torch.arange(window_size, device=device) + matrix = (base - window_size + 1).clamp(min=0) + offsets + matrix = torch.where(matrix > base, -1, matrix) + return matrix.unsqueeze(0).expand(batch_size, -1, -1) + + +def _get_compress_topk_idxs( + ratio: int, batch_size: int, seqlen: int, offset: int, device: torch.device +) -> torch.Tensor: + n_compressed = seqlen // ratio + matrix = torch.arange(n_compressed, device=device).repeat(seqlen, 1) + mask = matrix >= torch.arange(1, seqlen + 1, device=device).unsqueeze(1) // ratio + matrix = torch.where(mask, -1, matrix + offset) + return matrix.unsqueeze(0).expand(batch_size, -1, -1) + + +def _native_sparse_attn( + query: torch.Tensor, + kv_full: torch.Tensor, + attn_sink: torch.Tensor, + topk_indices: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + sq, batch_size, num_heads, head_dim = query.size() + kv_t = kv_full.permute(1, 0, 2) + chunk_size = 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) + 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_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) + + 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( + index_scores: torch.Tensor, + topk_indices: torch.Tensor, + query: torch.Tensor, + compressed_kv: torch.Tensor, + attn_sink: torch.Tensor, + softmax_scale: float, + loss_coeff: float, + sparse_loss: bool, + calculate_per_token_loss: bool, +) -> torch.Tensor: + batch_size, seqlen, _ = topk_indices.size() + num_heads, head_dim = query.size(2), query.size(3) + n_compressed = compressed_kv.size(0) + + sink = attn_sink.detach().view(1, num_heads, 1, 1).float() + q = query.detach().permute(1, 2, 0, 3).float() + compressed_kv_t = compressed_kv.detach().permute(1, 0, 2) + + if sparse_loss: + safe_indices = topk_indices.clamp(min=0).long() + valid = topk_indices >= 0 + row_valid = valid.any(dim=-1, keepdim=True) + + predict_logits = torch.gather(index_scores, dim=-1, index=safe_indices) + predict_logits = predict_logits.masked_fill(~valid, float("-inf")) + predict_logits = predict_logits.masked_fill(~row_valid, 0.0) + predict = F.softmax(predict_logits, dim=-1, dtype=torch.float32) + predict = predict * row_valid.float() + + selected_kv = torch.gather( + compressed_kv_t.unsqueeze(1).expand(-1, seqlen, -1, -1), + dim=2, + index=safe_indices.unsqueeze(-1).expand(-1, -1, -1, head_dim), + ) + attn_scores = torch.einsum("bhsd,bskd->bhsk", q, selected_kv.float()) + attn_scores = attn_scores * softmax_scale + attn_scores = attn_scores.masked_fill(~valid.unsqueeze(1), float("-inf")) + else: + # Dense loss: KL is computed over the FULL compressed-KV axis (not + # just topk). Index-side and attention-side both use the kernel's + # ratio-causal mask, which we derive analytically from the + # compress_ratio (= seqlen / n_compressed): position k of the + # compressed-KV is valid for query row q iff k < (q + 1) // ratio. + compress_ratio = seqlen // n_compressed + k_idx = torch.arange(n_compressed, device=index_scores.device) + valid_per_q = ( + torch.arange(1, seqlen + 1, device=index_scores.device) // compress_ratio + ).clamp(max=n_compressed) + finite_pos = k_idx.view(1, 1, -1) < valid_per_q.view(1, -1, 1) # (1, sq, n_compressed) + finite_pos = finite_pos.expand(batch_size, -1, -1) + row_valid = finite_pos.any(dim=-1, keepdim=True) + + predict_logits = index_scores.masked_fill(~finite_pos, float("-inf")) + predict_logits = predict_logits.masked_fill(~row_valid, 0.0) + predict = F.softmax(predict_logits, dim=-1, dtype=torch.float32) + predict = predict * row_valid.float() + + attn_scores = torch.einsum("bhsd,bkd->bhsk", q, compressed_kv_t.float()) + attn_scores = attn_scores * softmax_scale + attn_mask = finite_pos.unsqueeze(1).expand(-1, num_heads, -1, -1) + attn_scores = attn_scores.masked_fill(~attn_mask, float("-inf")) + + score_max = torch.max(attn_scores.max(dim=-1, keepdim=True).values, sink) + exp_scores = torch.exp(attn_scores - score_max) + exp_sink = torch.exp(sink - score_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)) + loss = kl_per_row.sum() if calculate_per_token_loss else kl_per_row.mean() + return loss_coeff * loss + + +def _native_unfused_sparse_indexer_loss( + index_scores: torch.Tensor, + topk_indices: torch.Tensor, + query: torch.Tensor, + compressed_kv: torch.Tensor, + softmax_scale: float, + loss_coeff: float, + sparse_loss: bool, + causal_mask: torch.Tensor, + calculate_per_token_loss: bool, +) -> torch.Tensor: + sq, batch_size, num_heads, _ = query.size() + sk = compressed_kv.size(0) + mask = causal_mask.to(dtype=torch.float32) + + attention_scores = torch.einsum( + "sbhd,tbd->bhst", query.detach().float(), compressed_kv.detach().float() + ) + attention_scores = attention_scores * softmax_scale + attention_scores = attention_scores + mask.view(batch_size, 1, sq, sk) + index_scores = index_scores + mask + + if sparse_loss: + index_mask = torch.full( + (batch_size, sq, sk), float("-inf"), dtype=torch.float32, device=index_scores.device + ).scatter_(-1, topk_indices.clamp(min=0), 0) + attention_scores = attention_scores + index_mask.view(batch_size, 1, sq, sk) + index_scores = index_scores + index_mask + + row_valid = (mask > float("-inf")).any(dim=-1) + attn_row_mask = row_valid.view(batch_size, 1, sq, 1) + idx_row_mask = row_valid.view(batch_size, sq, 1) + + attention_scores = attention_scores.masked_fill(~attn_row_mask, 0.0) + index_scores = index_scores.masked_fill(~idx_row_mask, 0.0) + + attention_probs = F.softmax(attention_scores, dim=-1, dtype=torch.float32) + predict = F.softmax(index_scores, dim=-1, dtype=torch.float32) + attention_probs = attention_probs * attn_row_mask.float() + predict = predict * idx_row_mask.float() + + target = attention_probs.sum(dim=1) + target = target / target.sum(dim=-1, keepdim=True) + 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, kl_per_row, torch.zeros_like(kl_per_row)) + loss = kl_per_row.sum() if calculate_per_token_loss else kl_per_row.mean() + return loss * loss_coeff + + +class NativeCompressor(nn.Module): + def __init__( + self, config: MLATransformerConfig, compress_ratio: int, head_dim: int, rotate: bool + ): + super().__init__() + self.compress_ratio = compress_ratio + self.head_dim = head_dim + self.overlap = compress_ratio == 4 + self.coff = 1 + int(self.overlap) + self.rotate = rotate + self.qk_pos_emb_head_dim = config.qk_pos_emb_head_dim + self.rope_base = ( + config.csa_compress_rotary_base if compress_ratio > 1 else config.rotary_base + ) + # YaRN frequency interpolation is enabled only for compressed sequences + # (matches ``DSv4HybridAttention``'s ``use_compressed_yarn = ratio > 1``). + if compress_ratio > 1: + self._rope_yarn_kwargs = dict( + original_seq_len=config.original_max_position_embeddings, + factor=config.rotary_scaling_factor, + beta_fast=config.beta_fast, + beta_slow=config.beta_slow, + ) + else: + self._rope_yarn_kwargs = dict() + + self.linear_wkv = nn.Linear(config.hidden_size, self.coff * head_dim, bias=False) + self.linear_wgate = nn.Linear(config.hidden_size, self.coff * head_dim, bias=False) + self.ape = nn.Parameter( + torch.empty(compress_ratio, self.coff * head_dim, dtype=torch.float32) + ) + self.norm = nn.RMSNorm(head_dim, eps=config.layernorm_epsilon) + + def _overlap_transform(self, tensor: torch.Tensor, fill_value: float = 0) -> torch.Tensor: + n_groups, ratio, batch_size, _ = tensor.size() + new_tensor = tensor.new_full((n_groups, 2 * ratio, batch_size, self.head_dim), fill_value) + new_tensor[:, ratio:] = tensor[:, :, :, self.head_dim :] + new_tensor[1:, :ratio] = tensor[:-1, :, :, : self.head_dim] + return new_tensor + + def forward(self, x: torch.Tensor) -> torch.Tensor | None: + sq, batch_size, _ = x.size() + ratio = self.compress_ratio + if sq < ratio: + return None + + kv = self.linear_wkv(x) + score = self.linear_wgate(x) + + cutoff = (sq // ratio) * ratio + kv = kv[:cutoff] + score = score[:cutoff] + n_compressed = cutoff // ratio + + kv = kv.view(n_compressed, ratio, batch_size, -1) + score = score.view(n_compressed, ratio, batch_size, -1) + score = score + self.ape.view(1, ratio, 1, -1) + + if self.overlap: + kv = self._overlap_transform(kv, fill_value=0) + score = self._overlap_transform(score, fill_value=float("-inf")) + + kv = (kv * torch.softmax(score, dim=1)).sum(dim=1) + kv = self.norm(kv.to(x.dtype)) + + pos_dim = self.qk_pos_emb_head_dim + content, rotary = torch.split(kv, [self.head_dim - pos_dim, pos_dim], dim=-1) + freqs_cis = _precompute_freqs_cis( + pos_dim, + n_compressed * ratio, + device=x.device, + base=self.rope_base, + **self._rope_yarn_kwargs, + ) + freqs_cis = freqs_cis[: n_compressed * ratio : ratio][:n_compressed] + rotary = _apply_rotary_emb(rotary, freqs_cis) + kv = torch.cat([content, rotary], dim=-1) + + if self.rotate: + kv = _native_hadamard_transform(kv) + return kv + + +class NativeCSAIndexer(nn.Module): + def __init__(self, config: MLATransformerConfig, compress_ratio: int): + super().__init__() + self.compress_ratio = compress_ratio + self.index_n_heads = config.dsa_indexer_n_heads + self.index_head_dim = config.dsa_indexer_head_dim + self.index_topk = config.dsa_indexer_topk + self.qk_pos_emb_head_dim = config.qk_pos_emb_head_dim + self.softmax_scale = self.index_head_dim**-0.5 + self.apply_dsa_kernel_fusion = config.apply_dsa_kernel_fusion + self.rope_base = config.csa_compress_rotary_base + # CSA indexer is only instantiated for ``compress_ratio == 4``, which is + # always the YaRN-enabled branch on the production side. + self._rope_yarn_kwargs = dict( + original_seq_len=config.original_max_position_embeddings, + factor=config.rotary_scaling_factor, + beta_fast=config.beta_fast, + beta_slow=config.beta_slow, + ) + + self.linear_wq_b = nn.Linear( + config.q_lora_rank, self.index_n_heads * self.index_head_dim, bias=False + ) + self.linear_weights_proj = nn.Linear(config.hidden_size, self.index_n_heads, bias=False) + self.compressor = NativeCompressor( + config=config, compress_ratio=compress_ratio, head_dim=self.index_head_dim, rotate=True + ) + + def forward_before_topk( + self, x: torch.Tensor, qr: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + sq, batch_size, _ = x.size() + q = self.linear_wq_b(qr).view(sq, batch_size, self.index_n_heads, self.index_head_dim) + pos_dim = self.qk_pos_emb_head_dim + q_content, q_rotary = torch.split(q, [self.index_head_dim - pos_dim, pos_dim], dim=-1) + freqs_cis = _precompute_freqs_cis( + pos_dim, sq, device=x.device, base=self.rope_base, **self._rope_yarn_kwargs + ) + q_rotary = _apply_rotary_emb(q_rotary, freqs_cis) + q = _native_hadamard_transform(torch.cat([q_content, q_rotary], dim=-1)) + + k = self.compressor(x) + weights = self.linear_weights_proj(x) * (self.index_n_heads**-0.5) + return q, k, weights + + def forward( + self, x: torch.Tensor, qr: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, weights = self.forward_before_topk(x, qr) + weights_scaled = weights.float() * self.softmax_scale + if self.apply_dsa_kernel_fusion: + weights_scaled = weights_scaled.to(weights.dtype).float() + scores = torch.einsum("sbhd,tbd->sbht", q.float(), k.float()) + scores = torch.relu(scores) * weights_scaled.unsqueeze(-1) + scores = scores.sum(dim=2).transpose(0, 1) + + sq = x.size(0) + n_compressed = k.size(0) + valid_per_query = ( + torch.arange(1, sq + 1, device=x.device).unsqueeze(0) // self.compress_ratio + ).clamp(max=n_compressed) + invalid = torch.arange(n_compressed, device=x.device).view( + 1, 1, -1 + ) >= valid_per_query.unsqueeze(-1) + scores = scores.masked_fill(invalid.expand_as(scores), float("-inf")) + + topk = min(self.index_topk, n_compressed) + topk_scores, topk_indices = scores.topk(topk, dim=-1) + topk_indices = torch.where(topk_scores.isneginf(), -1, topk_indices) + return q, k, weights, scores, topk_indices + + +class NativeCompressedSparseAttention(nn.Module): + def __init__(self, config: MLATransformerConfig, compress_ratio: int): + super().__init__() + self.compress_ratio = compress_ratio + self.window_size = config.csa_window_size + self.num_heads = config.num_attention_heads + self.head_dim = config.v_head_dim + self.softmax_scale = self.head_dim**-0.5 + self.indexer_loss_coeff = config.dsa_indexer_loss_coeff + self.indexer_use_sparse_loss = config.dsa_indexer_use_sparse_loss + self.calculate_per_token_loss = config.calculate_per_token_loss + self.apply_dsa_kernel_fusion = config.apply_dsa_kernel_fusion + + self.attn_sink = nn.Parameter(torch.zeros(self.num_heads, dtype=torch.float32)) + self.compressor = ( + NativeCompressor( + config=config, compress_ratio=compress_ratio, head_dim=self.head_dim, rotate=False + ) + if compress_ratio > 1 + else None + ) + self.indexer = NativeCSAIndexer(config, compress_ratio) if compress_ratio == 4 else None + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + x: torch.Tensor, + qr: torch.Tensor, + pg_collection, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + sq, batch_size, _, _ = query.size() + kv = key.squeeze(-2) + n_compressed = 0 + + if self.compressor is not None: + compressed_kv = self.compressor(x) + if compressed_kv is not None: + kv_full = torch.cat([kv, compressed_kv], dim=0) + n_compressed = compressed_kv.size(0) + else: + kv_full = kv + else: + compressed_kv = None + kv_full = kv + + window_idxs = _get_window_topk_idxs(self.window_size, batch_size, sq, query.device) + indexer_loss = None + if self.compress_ratio > 1 and n_compressed > 0: + offset = sq + if self.indexer is not None: + q_idx, k_idx, weights_idx, index_scores, topk_compressed = self.indexer( + x.detach(), qr.detach() + ) + topk_compressed_for_attn = torch.where( + topk_compressed >= 0, topk_compressed + offset, -1 + ) + + if not self.apply_dsa_kernel_fusion: + causal_mask = ( + torch.arange(n_compressed, device=x.device).unsqueeze(0).expand(sq, -1) + ) + positions = torch.arange(1, sq + 1, device=x.device).unsqueeze(1) + causal_mask = ( + torch.where( + causal_mask >= positions // self.compress_ratio, float("-inf"), 0.0 + ) + .unsqueeze(0) + .expand(batch_size, -1, -1) + ) + indexer_loss = _native_unfused_sparse_indexer_loss( + index_scores, + topk_compressed, + query.detach(), + compressed_kv.detach(), + self.softmax_scale, + self.indexer_loss_coeff, + self.indexer_use_sparse_loss, + causal_mask, + self.calculate_per_token_loss, + ) + else: + indexer_loss = _native_fused_sparse_indexer_loss( + index_scores, + topk_compressed, + query, + compressed_kv, + self.attn_sink, + self.softmax_scale, + self.indexer_loss_coeff, + self.indexer_use_sparse_loss, + self.calculate_per_token_loss, + ) + else: + topk_compressed_for_attn = _get_compress_topk_idxs( + self.compress_ratio, batch_size, sq, offset, query.device + ) + if self.indexer is not None and self.apply_dsa_kernel_fusion: + topk_idxs = torch.cat([topk_compressed_for_attn, window_idxs], dim=-1) + else: + topk_idxs = torch.cat([window_idxs, topk_compressed_for_attn], dim=-1) + else: + topk_idxs = window_idxs + + output = _native_sparse_attn(query, kv_full, self.attn_sink, topk_idxs, self.softmax_scale) + return output, indexer_loss + + +class NativeDSv4HybridAttention(nn.Module): + def __init__(self, config: MLATransformerConfig, compress_ratio: int): + super().__init__() + self.config = config + self.compress_ratio = compress_ratio + self.num_heads = config.num_attention_heads + self.head_dim = config.v_head_dim + self.pos_dim = config.qk_pos_emb_head_dim + self.nope_dim = config.v_head_dim - config.qk_pos_emb_head_dim + self.rope_base = ( + config.csa_compress_rotary_base if compress_ratio > 1 else config.rotary_base + ) + if compress_ratio > 1: + self._rope_yarn_kwargs = dict( + original_seq_len=config.original_max_position_embeddings, + factor=config.rotary_scaling_factor, + beta_fast=config.beta_fast, + beta_slow=config.beta_slow, + ) + else: + self._rope_yarn_kwargs = dict() + + self.linear_q_down_proj = nn.Linear(config.hidden_size, config.q_lora_rank, bias=False) + self.q_layernorm = nn.RMSNorm(config.q_lora_rank, eps=config.layernorm_epsilon) + self.linear_q_up_proj = nn.Linear( + config.q_lora_rank, config.num_attention_heads * config.v_head_dim, bias=False + ) + self.linear_kv_proj = nn.Linear(config.hidden_size, config.v_head_dim, bias=False) + self.kv_layernorm = nn.RMSNorm(config.v_head_dim, eps=config.layernorm_epsilon) + self.core_attention = NativeCompressedSparseAttention(config, compress_ratio) + group_in = (config.num_attention_heads * config.v_head_dim) // config.o_groups + self.linear_o_group_proj = nn.Parameter( + torch.empty(config.o_groups * config.o_lora_rank, group_in) + ) + self.linear_proj = nn.Linear( + config.o_groups * config.o_lora_rank, config.hidden_size, bias=False + ) + + def forward( + self, hidden_states: torch.Tensor, pg_collection + ) -> tuple[torch.Tensor, torch.Tensor | None]: + sq, batch_size, _ = hidden_states.size() + freqs_cis = _precompute_freqs_cis( + self.pos_dim, sq, hidden_states.device, self.rope_base, **self._rope_yarn_kwargs + ) + + qr = self.q_layernorm(self.linear_q_down_proj(hidden_states)) + query = self.linear_q_up_proj(qr).view(sq, batch_size, self.num_heads, self.head_dim) + query = _native_q_rms_norm(query, self.config.layernorm_epsilon) + q_content, q_rotary = torch.split(query, [self.nope_dim, self.pos_dim], dim=-1) + query = torch.cat([q_content, _apply_rotary_emb(q_rotary, freqs_cis)], dim=-1) + + key = self.kv_layernorm(self.linear_kv_proj(hidden_states)) + k_content, k_rotary = torch.split(key, [self.nope_dim, self.pos_dim], dim=-1) + key = torch.cat([k_content, _apply_rotary_emb(k_rotary, freqs_cis)], dim=-1) + key = key.unsqueeze(-2) + + core_out, indexer_loss = self.core_attention( + query=query, key=key, x=hidden_states, qr=qr, pg_collection=pg_collection + ) + + core_out = core_out.view(sq, batch_size, self.num_heads, self.head_dim) + out_content, out_rotary = torch.split(core_out, [self.nope_dim, self.pos_dim], dim=-1) + core_out = torch.cat( + [out_content, _apply_rotary_emb(out_rotary, freqs_cis, inverse=True)], dim=-1 + ) + core_out = core_out.view(sq, batch_size, -1) + + core_out = core_out.view(sq, batch_size, self.config.o_groups, -1) + wo_a = self.linear_o_group_proj.view(self.config.o_groups, self.config.o_lora_rank, -1) + core_out = torch.einsum("...gd,grd->...gr", core_out, wo_a) + core_out = core_out.reshape(sq, batch_size, -1) + return self.linear_proj(core_out), indexer_loss + + +def _cosine_sim(a: torch.Tensor, b: torch.Tensor) -> 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: + 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_similarity(a: torch.Tensor, b: torch.Tensor, label: str, eps: float): + assert torch.isfinite(a).all() + assert torch.isfinite(b).all() + cosine_sim = _cosine_sim(a, b) + tensor_sim = _tensor_sim(a, b) + assert cosine_sim > 1 - eps, f"{label}: cosine_sim={cosine_sim:.10f}, eps={eps}" + assert tensor_sim > 1 - eps, f"{label}: tensor_sim={tensor_sim:.10f}, eps={eps}" + + +def _copy_real_params_to_native(real_layer: nn.Module, native_layer: nn.Module): + real_params = dict(real_layer.named_parameters()) + for name, native_param in native_layer.named_parameters(): + assert name in real_params, f"Missing real parameter for native parameter {name}" + real_param = real_params[name] + assert ( + native_param.shape == real_param.shape + ), f"Shape mismatch for {name}: native={native_param.shape}, real={real_param.shape}" + native_param.data = real_param.data.to( + device=native_param.device, dtype=real_param.dtype + ).clone() + return real_params + + +def _make_thd_packed_seq_params(seg_lens, device='cuda'): + """Build ``PackedSeqParams(qkv_format='thd', ...)`` for self-attention + from a list of per-segment lengths. + """ + cu_seqlens = torch.tensor( + [0] + list(torch.tensor(seg_lens, dtype=torch.int64).cumsum(0).tolist()), + dtype=torch.int32, + device=device, + ) + max_len = int(max(seg_lens)) if seg_lens else 0 + return PackedSeqParams( + cu_seqlens_q=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_kv_padded=cu_seqlens, + max_seqlen_q=max_len, + max_seqlen_kv=max_len, + qkv_format='thd', + ) + + +def _skip_if_real_kernels_unavailable(*, sm_min: int = 9, need_flash_mla: bool = False): + """Pytest-side gate for real-kernel tests. Raises ``pytest.skip`` if + any of the runtime dependencies are missing. + """ + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + sm_major = torch.cuda.get_device_capability()[0] + if sm_major < sm_min: + pytest.skip(f"requires SM{sm_min}+, found SM{sm_major}") + cudnn = pytest.importorskip("cudnn") + from packaging.version import Version + + if Version(cudnn.__version__) < Version("1.24.0"): + pytest.skip(f"requires cudnn>=1.24.0, found {cudnn.__version__}") + if not hasattr(cudnn, 'DSA'): + pytest.skip("cudnn.DSA namespace not available") + if need_flash_mla: + pytest.importorskip("flash_mla") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridNativeParity: + + @classmethod + def setup_class(cls): + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=1) + + @classmethod + def teardown_class(cls): + Utils.destroy_model_parallel() + + def setup_method(self): + DSAIndexerLossAutoScaler.main_loss_backward_scale = None + torch.manual_seed(_SEED) + torch.cuda.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + def teardown_method(self): + gc.collect() + torch.cuda.empty_cache() + + @pytest.mark.parametrize(("backend", "apply_dsa_kernel_fusion"), _DSA_BACKENDS) + @pytest.mark.parametrize("variant", ["flash"]) + @pytest.mark.parametrize("compress_ratio", [1, 4, 128]) + @pytest.mark.parametrize( + ("seqlen", "calculate_per_token_loss", "dsa_indexer_use_sparse_loss"), + [(512, True, True)], + ) + def test_attention_matches_native_reference( + self, + variant: str, + compress_ratio: int, + seqlen: int, + backend: str, + apply_dsa_kernel_fusion: bool, + calculate_per_token_loss: bool, + dsa_indexer_use_sparse_loss: bool, + ): + if apply_dsa_kernel_fusion: + _skip_if_real_kernels_unavailable() + major, _ = torch.cuda.get_device_capability() + if ( + major == 9 + and apply_dsa_kernel_fusion + and compress_ratio == 4 + and not dsa_indexer_use_sparse_loss + ): + pytest.skip("cuDNN Frontend SM90 dense DSA is not supported") + + config = _make_config( + variant, + compress_ratio, + apply_dsa_kernel_fusion=apply_dsa_kernel_fusion, + calculate_per_token_loss=calculate_per_token_loss, + dsa_indexer_use_sparse_loss=dsa_indexer_use_sparse_loss, + ) + fwd_eps = ( + _FUSED_FWD_SIMILARITY_EPS if apply_dsa_kernel_fusion else _UNFUSED_FWD_SIMILARITY_EPS + ) + bwd_eps = ( + _FUSED_BWD_SIMILARITY_EPS if apply_dsa_kernel_fusion else _UNFUSED_BWD_SIMILARITY_EPS + ) + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=["tp", "cp"]) + spec = get_dsv4_hybrid_module_spec_for_backend(config=config, backend=TESpecProvider()) + + mcore_ratio = 0 if compress_ratio == 1 else compress_ratio + real_layer = build_module( + spec, config=config, layer_number=1, cp_comm_type=None, pg_collection=pg_collection + ).cuda() + native_layer = NativeDSv4HybridAttention(config, mcore_ratio).cuda() + real_params = _copy_real_params_to_native(real_layer, native_layer) + + bsz = 1 + for _ in range(1): + hidden_states = torch.randn( + seqlen, + bsz, + config.hidden_size, + dtype=torch.bfloat16, + device="cuda", + requires_grad=True, + ) + hidden_states_native = hidden_states.detach().clone().requires_grad_(True) + grad = torch.randn_like(hidden_states) + + real_out, _ = real_layer(hidden_states=hidden_states, attention_mask=None) + native_out, native_indexer_loss = native_layer(hidden_states_native, pg_collection) + + _assert_similarity( + real_out.detach(), + native_out.detach(), + f"{backend}-{variant}-{compress_ratio}-{seqlen}:out", + eps=fwd_eps, + ) + + real_out.backward(grad) + native_out.backward(grad) + if native_indexer_loss is not None: + native_indexer_loss.backward() + + _assert_similarity( + hidden_states.grad, + hidden_states_native.grad, + f"{backend}-{variant}-{compress_ratio}-{seqlen}:hidden_grad", + eps=bwd_eps, + ) + + for name, native_param in native_layer.named_parameters(): + real_param = real_params[name] + if compress_ratio != 4 and ".indexer." in name: + continue + assert native_param.grad is not None, f"Missing native grad for {name}" + assert real_param.grad is not None, f"Missing real grad for {name}" + _assert_similarity( + real_param.grad, + native_param.grad, + f"{backend}-{variant}-{compress_ratio}-{seqlen}:param_grad:{name}", + eps=bwd_eps, + ) + + del real_layer, native_layer, real_params + del hidden_states, hidden_states_native, real_out, native_out, grad + if native_indexer_loss is not None: + del native_indexer_loss + gc.collect() + torch.cuda.empty_cache() + + @pytest.mark.parametrize(("backend", "apply_dsa_kernel_fusion"), _DSA_BACKENDS) + @pytest.mark.parametrize("variant", ["flash"]) + @pytest.mark.parametrize("compress_ratio", [1, 4, 128]) + @pytest.mark.parametrize( + ("seqlen", "dsa_indexer_use_sparse_loss"), + [(512, False), (512, True)], + ) + def test_thd_attention_matches_native_reference( + self, + variant: str, + compress_ratio: int, + seqlen: int, + backend: str, + apply_dsa_kernel_fusion: bool, + dsa_indexer_use_sparse_loss: bool, + ): + """THD (packed-sequence) variant of test_attention_matches_native_reference. + + Runs the real layer with a single-segment THD packed_seq_params + (equivalent to SBHD B=1) and compares forward output and backward + gradients against the native reference. + """ + if apply_dsa_kernel_fusion: + _skip_if_real_kernels_unavailable() + major, _ = torch.cuda.get_device_capability() + if ( + major == 9 + and apply_dsa_kernel_fusion + and compress_ratio == 4 + and not dsa_indexer_use_sparse_loss + ): + pytest.skip("cuDNN Frontend SM90 THD dense DSA has cache and stream bugs") + + config = _make_config( + variant, + compress_ratio, + apply_dsa_kernel_fusion=apply_dsa_kernel_fusion, + calculate_per_token_loss=True, + dsa_indexer_use_sparse_loss=dsa_indexer_use_sparse_loss, + ) + fwd_eps = ( + _FUSED_FWD_SIMILARITY_EPS if apply_dsa_kernel_fusion else _UNFUSED_FWD_SIMILARITY_EPS + ) + bwd_eps = ( + _FUSED_BWD_SIMILARITY_EPS if apply_dsa_kernel_fusion else _UNFUSED_BWD_SIMILARITY_EPS + ) + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=["tp", "cp"]) + spec = get_dsv4_hybrid_module_spec_for_backend(config=config, backend=TESpecProvider()) + + mcore_ratio = 0 if compress_ratio == 1 else compress_ratio + real_layer = build_module( + spec, config=config, layer_number=1, cp_comm_type=None, pg_collection=pg_collection + ).cuda() + native_layer = NativeDSv4HybridAttention(config, mcore_ratio).cuda() + real_params = _copy_real_params_to_native(real_layer, native_layer) + + bsz = 1 + for _ in range(1): + hidden_states = torch.randn( + seqlen, + bsz, + config.hidden_size, + dtype=torch.bfloat16, + device="cuda", + requires_grad=True, + ) + hidden_states_native = hidden_states.detach().clone().requires_grad_(True) + grad = torch.randn_like(hidden_states) + + packed = _make_thd_packed_seq_params([seqlen]) + real_out, _ = real_layer( + hidden_states=hidden_states, attention_mask=None, packed_seq_params=packed + ) + native_out, native_indexer_loss = native_layer(hidden_states_native, pg_collection) + + _assert_similarity( + real_out.detach(), + native_out.detach(), + f"thd-{backend}-{variant}-{compress_ratio}-{seqlen}:out", + eps=fwd_eps, + ) + + real_out.backward(grad) + native_out.backward(grad) + if native_indexer_loss is not None: + native_indexer_loss.backward() + + _assert_similarity( + hidden_states.grad, + hidden_states_native.grad, + f"thd-{backend}-{variant}-{compress_ratio}-{seqlen}:hidden_grad", + eps=bwd_eps, + ) + + for name, native_param in native_layer.named_parameters(): + real_param = real_params[name] + if compress_ratio != 4 and ".indexer." in name: + continue + assert native_param.grad is not None, f"Missing native grad for {name}" + assert real_param.grad is not None, f"Missing real grad for {name}" + _assert_similarity( + real_param.grad, + native_param.grad, + f"thd-{backend}-{variant}-{compress_ratio}-{seqlen}:param_grad:{name}", + eps=bwd_eps, + ) + + del real_layer, native_layer, real_params + del hidden_states, hidden_states_native, real_out, native_out, grad, packed + if native_indexer_loss is not None: + del native_indexer_loss + gc.collect() + torch.cuda.empty_cache() + + @pytest.mark.parametrize(("backend", "apply_dsa_kernel_fusion"), _DSA_BACKENDS) + @pytest.mark.parametrize("variant", ["flash"]) + @pytest.mark.parametrize("compress_ratio", [1, 4, 128]) + @pytest.mark.parametrize( + ("seg_lens", "dsa_indexer_use_sparse_loss"), + [ + # pytest.param([152, 1024, 2345], False, id="three-seg-dense"), + pytest.param([152, 1024, 2345], True, id="three-seg-sparse") + ], + ) + def test_thd_multiseg_attention_matches_native_reference( + self, + variant: str, + compress_ratio: int, + seg_lens: list, + backend: str, + apply_dsa_kernel_fusion: bool, + dsa_indexer_use_sparse_loss: bool, + ): + """Multi-segment THD parity against per-segment native references. + + The single-segment ``test_thd_attention_matches_native_reference`` + cannot distinguish per-segment RoPE striding from global striding: + with one segment starting at offset 0 the two coincide bit-for-bit. + Real packed sequences reset RoPE positions *per segment* (the kernel + indexes the globally-strided cos/sin table via ``cu_seqlens``), so the + correct oracle is the native reference run **independently per + segment** — each segment seeing positions ``0..seg_len-1`` — with the + outputs concatenated. Comparing the packed real layer against that + oracle exercises cross-segment RoPE / compression positioning, the + class of bug that single-segment and padding-invariance tests miss. + + Segment lengths are multiples of 128 (== ``csa_window_size`` and the + max compress ratio) so compression is exact at every ratio. + """ + if apply_dsa_kernel_fusion: + _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: + pytest.skip("seqlen > 4096 may OOM on Hopper with unfused DSA implementation") + + config = _make_config( + variant, + compress_ratio, + apply_dsa_kernel_fusion=apply_dsa_kernel_fusion, + calculate_per_token_loss=True, + dsa_indexer_use_sparse_loss=dsa_indexer_use_sparse_loss, + ) + fwd_eps = ( + _FUSED_FWD_SIMILARITY_EPS if apply_dsa_kernel_fusion else _UNFUSED_FWD_SIMILARITY_EPS + ) + bwd_eps = ( + _FUSED_BWD_SIMILARITY_EPS if apply_dsa_kernel_fusion else _UNFUSED_BWD_SIMILARITY_EPS + ) + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=["tp", "cp"]) + spec = get_dsv4_hybrid_module_spec_for_backend(config=config, backend=TESpecProvider()) + + mcore_ratio = 0 if compress_ratio == 1 else compress_ratio + real_layer = build_module( + spec, config=config, layer_number=1, cp_comm_type=None, pg_collection=pg_collection + ).cuda() + native_layer = NativeDSv4HybridAttention(config, mcore_ratio).cuda() + real_params = _copy_real_params_to_native(real_layer, native_layer) + + hidden_states = torch.randn( + total_T, 1, config.hidden_size, dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + hidden_states_native = hidden_states.detach().clone().requires_grad_(True) + grad = torch.randn_like(hidden_states) + + # ---- Real packed-sequence (THD) run ---------------------------------- + packed = _make_thd_packed_seq_params(seg_lens) + real_out, _ = real_layer( + hidden_states=hidden_states, attention_mask=None, packed_seq_params=packed + ) + + # ---- Native oracle: each segment as an independent B=1 sequence ------ + # Slicing the single ``hidden_states_native`` leaf keeps every segment's + # input grad flowing back into one tensor (comparable to the real + # layer's packed grad); reusing one ``native_layer`` accumulates param + # grads across segments exactly as the packed real layer does. + seg_label = "_".join(map(str, seg_lens)) + seg_outs = [] + seg_losses = [] + start = 0 + for seg_len in seg_lens: + seg_in = hidden_states_native[start : start + seg_len] + seg_out, seg_loss = native_layer(seg_in, pg_collection) + seg_outs.append(seg_out) + if seg_loss is not None: + seg_losses.append(seg_loss) + start += seg_len + native_out = torch.cat(seg_outs, dim=0) + + _assert_similarity( + real_out.detach(), + native_out.detach(), + f"thd-multiseg-{backend}-{variant}-{compress_ratio}-{seg_label}:out", + eps=fwd_eps, + ) + + real_out.backward(grad) + native_out.backward(grad) + if seg_losses: + # per_token_loss=True => each segment's loss is a row-sum; summing + # across segments equals the packed layer's whole-sequence sum. + torch.stack(seg_losses).sum().backward() + + _assert_similarity( + hidden_states.grad, + hidden_states_native.grad, + f"thd-multiseg-{backend}-{variant}-{compress_ratio}-{seg_label}:hidden_grad", + eps=bwd_eps, + ) + + for name, native_param in native_layer.named_parameters(): + real_param = real_params[name] + if compress_ratio != 4 and ".indexer." in name: + continue + assert native_param.grad is not None, f"Missing native grad for {name}" + assert real_param.grad is not None, f"Missing real grad for {name}" + _assert_similarity( + real_param.grad, + native_param.grad, + f"thd-multiseg-{backend}-{variant}-{compress_ratio}-{seg_label}:param_grad:{name}", + eps=bwd_eps, + ) + + del real_layer, native_layer, real_params + del hidden_states, hidden_states_native, real_out, native_out, grad, packed + del seg_outs, seg_losses + gc.collect() + torch.cuda.empty_cache() + + @pytest.mark.parametrize(("backend", "apply_dsa_kernel_fusion"), _DSA_BACKENDS) + @pytest.mark.parametrize("variant", ["flash"]) + @pytest.mark.parametrize("compress_ratio", [4, 128]) + @pytest.mark.parametrize("dsa_indexer_use_sparse_loss", [True, False]) + @pytest.mark.parametrize( + ("seg_lens", "pad_max_seqlen", "pad_max_num_seqs"), + [ + pytest.param([512], 640, 4, id="single-seg-padded"), + pytest.param([256, 256], 640, 4, id="two-seg-padded"), + pytest.param([200, 150, 912], 2048, 8, id="three-seg-padded"), + ], + ) + def test_thd_padded_attention_matches_unpadded( + self, + variant: str, + compress_ratio: int, + seg_lens: list, + pad_max_seqlen: int, + pad_max_num_seqs: int, + backend: str, + dsa_indexer_use_sparse_loss: bool, + apply_dsa_kernel_fusion: bool, + ): + """Verify that THD padding does not corrupt real tokens' output. + + Runs the same real layer twice — once with padding (static shapes) + and once without — then asserts the forward output and backward + gradients for the real (non-padding) token positions are identical + within tolerance. + """ + if apply_dsa_kernel_fusion: + _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" + assert len(seg_lens) <= pad_max_num_seqs, "seg count must fit within pad_max_num_seqs" + + config = _make_config( + variant, + compress_ratio, + apply_dsa_kernel_fusion=apply_dsa_kernel_fusion, + calculate_per_token_loss=True, + dsa_indexer_use_sparse_loss=dsa_indexer_use_sparse_loss, + ) + fwd_eps = ( + _FUSED_FWD_SIMILARITY_EPS if apply_dsa_kernel_fusion else _UNFUSED_FWD_SIMILARITY_EPS + ) + bwd_eps = ( + _FUSED_BWD_SIMILARITY_EPS if apply_dsa_kernel_fusion else _UNFUSED_BWD_SIMILARITY_EPS + ) + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=["tp", "cp"]) + spec = get_dsv4_hybrid_module_spec_for_backend(config=config, backend=TESpecProvider()) + + real_layer = build_module( + spec, config=config, layer_number=1, cp_comm_type=None, pg_collection=pg_collection + ).cuda() + + hidden_states = torch.randn( + actual_T, 1, config.hidden_size, dtype=torch.bfloat16, device="cuda" + ) + + # ---- Unpadded run (reference) ---------------------------------------- + hidden_unpadded = hidden_states.detach().clone().requires_grad_(True) + packed_unpadded = _make_thd_packed_seq_params(seg_lens) + out_unpadded, _ = real_layer( + hidden_states=hidden_unpadded, attention_mask=None, packed_seq_params=packed_unpadded + ) + grad_unpadded = torch.randn_like(out_unpadded) + out_unpadded.backward(grad_unpadded) + + # ---- Padded run ------------------------------------------------------ + # Per-sequence padding: each segment is padded to the next multiple + # of compress_ratio, then the total is extended to pad_max_seqlen + # with a dummy tail segment. This matches the real data_schedule + # path where each sequence is individually padded to alignment. + from megatron.core.packed_seq_params import _pad_cu_seqlens + + align = max(compress_ratio, 4) + padded_seg_lens = [((sl + align - 1) // align) * align for sl in seg_lens] + padded_actual_T = sum(padded_seg_lens) + total_padded_T = pad_max_seqlen + assert padded_actual_T <= total_padded_T + + # Build padded hidden_states with per-segment intra-padding + tail. + hidden_padded = torch.zeros( + total_padded_T, 1, config.hidden_size, dtype=torch.bfloat16, device="cuda" + ) + real_offset = 0 + pad_offset = 0 + for rl, pl in zip(seg_lens, padded_seg_lens): + hidden_padded[pad_offset : pad_offset + rl] = hidden_states[ + real_offset : real_offset + rl + ] + real_offset += rl + pad_offset += pl + hidden_padded = hidden_padded.clone().requires_grad_(True) + + # cu_seqlens_q: unpadded real boundaries (cumsum of real lengths + # within the padded physical layout). + cu_seqlens_q_vals = [0] + offset = 0 + for rl, pl in zip(seg_lens, padded_seg_lens): + cu_seqlens_q_vals.append(offset + rl) + offset += pl + cu_seqlens_unpadded = torch.tensor(cu_seqlens_q_vals, dtype=torch.int32, device='cuda') + + # cu_seqlens_q_padded: padded boundaries (cumsum of padded lengths + # + dummy tail segment to total_padded_T). + padded_boundaries = [0] + for pl in padded_seg_lens: + padded_boundaries.append(padded_boundaries[-1] + pl) + if padded_actual_T < total_padded_T: + padded_boundaries.append(total_padded_T) + cu_seqlens_padded_raw = torch.tensor(padded_boundaries, dtype=torch.int32, device='cuda') + # Also extend unpadded with a zero-length dummy for the tail. + if padded_actual_T < total_padded_T: + cu_seqlens_unpadded = torch.cat( + [ + cu_seqlens_unpadded, + cu_seqlens_unpadded[-1:], # repeat last (real total unchanged) + ] + ) + + target_cu = pad_max_num_seqs + 1 + cu_seqlens_unpadded = _pad_cu_seqlens(cu_seqlens_unpadded, target_cu) + cu_seqlens_padded = _pad_cu_seqlens(cu_seqlens_padded_raw, target_cu) + max_padded_seg = max(padded_seg_lens + [total_padded_T - padded_actual_T]) + + packed_padded = PackedSeqParams( + qkv_format='thd', + cu_seqlens_q=cu_seqlens_unpadded, + cu_seqlens_kv=cu_seqlens_unpadded, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=max_padded_seg, + max_seqlen_kv=max_padded_seg, + ) + + out_padded, _ = real_layer( + hidden_states=hidden_padded, attention_mask=None, packed_seq_params=packed_padded + ) + # Build grad for padded buffer: scatter unpadded grad into real positions. + grad_padded = torch.zeros_like(out_padded) + real_offset = 0 + pad_offset = 0 + for rl, pl in zip(seg_lens, padded_seg_lens): + grad_padded[pad_offset : pad_offset + rl] = grad_unpadded[ + real_offset : real_offset + rl + ] + real_offset += rl + pad_offset += pl + out_padded.backward(grad_padded) + + # ---- Assertions: real tokens must match ------------------------------ + # Gather real-token positions from the padded output/grad. + real_positions = [] + pad_offset = 0 + for rl, pl in zip(seg_lens, padded_seg_lens): + real_positions.extend(range(pad_offset, pad_offset + rl)) + pad_offset += pl + real_positions = torch.tensor(real_positions, dtype=torch.long, device='cuda') + + label = f"thd-padded-{backend}-{variant}-r{compress_ratio}-segs{len(seg_lens)}" + _assert_similarity( + out_padded[real_positions].detach(), out_unpadded.detach(), f"{label}:out", eps=fwd_eps + ) + _assert_similarity( + hidden_padded.grad[real_positions], + hidden_unpadded.grad, + f"{label}:hidden_grad", + eps=bwd_eps, + ) + + del real_layer, hidden_states, hidden_unpadded, hidden_padded + del out_unpadded, out_padded, grad_unpadded, grad_padded + del packed_unpadded, packed_padded + gc.collect() + torch.cuda.empty_cache() diff --git a/tests/unit_tests/transformer/moe/test_aux_loss.py b/tests/unit_tests/transformer/moe/test_aux_loss.py index c8c7bf0dd0f..05b2f70a2a4 100644 --- a/tests/unit_tests/transformer/moe/test_aux_loss.py +++ b/tests/unit_tests/transformer/moe/test_aux_loss.py @@ -1,6 +1,8 @@ # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. import dataclasses +from types import SimpleNamespace +from unittest import mock import pytest import torch @@ -11,6 +13,7 @@ get_cuda_rng_tracker, model_parallel_cuda_manual_seed, ) +from megatron.core.transformer.moe import router as router_module from megatron.core.transformer.moe.moe_utils import ( clear_aux_losses_tracker, get_default_pg_collection, @@ -39,6 +42,50 @@ HAVE_ROUTER_FUSION = False +def test_per_token_aux_loss_scales_by_exact_reduced_valid_token_count(): + """Rank-skewed padding must use the reduced count, not local_count * group_size.""" + router = TopKRouter.__new__(TopKRouter) + torch.nn.Module.__init__(router) + router.config = SimpleNamespace( + mtp_num_layers=None, + mtp_use_repeated_layer=False, + num_layers=2, + ) + router.is_mtp_layer = False + router.layer_number = 1 + router.mtp_layer_number = None + router.calculate_per_token_loss = True + + activation = torch.ones(2) + aux_loss = torch.tensor(0.5) + reduce_group = mock.sentinel.reduce_group + + def add_remote_valid_tokens(token_count, group): + assert group is reduce_group + token_count.add_(3) + + tracker = mock.MagicMock() + with ( + mock.patch.object(torch.distributed, "all_reduce", side_effect=add_remote_valid_tokens), + mock.patch.object(router_module, "get_moe_metrics_tracker", return_value=tracker), + mock.patch.object( + router_module.MoEAuxLossAutoScaler, "apply", return_value=activation + ) as attach, + ): + result = router.attach_and_log_load_balancing_loss( + activation, + aux_loss_coeff=0.1, + aux_loss=aux_loss, + aux_loss_name="load_balancing_loss", + reduce_group=reduce_group, + valid_token_count=2, + aux_loss_scale_reduce_groups=(reduce_group,), + ) + + assert result is activation + torch.testing.assert_close(attach.call_args.args[1], torch.tensor(2.5)) + + class AuxlossTestContainer(MoEModelTestContainer): def partition_input(self, input): partitioned_input = input.chunk( diff --git a/tests/unit_tests/transformer/moe/test_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_grouped_mlp.py index 75fa941769a..33bd3553d91 100644 --- a/tests/unit_tests/transformer/moe/test_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_grouped_mlp.py @@ -38,6 +38,23 @@ def test_op_fuser_transformer_config_args_are_exposed(): assert args.moe_mlp_glu_interleave_size == 16 +def test_clamped_swiglu_allows_te_op_fuser(): + config = TransformerConfig( + num_layers=1, + hidden_size=16, + num_attention_heads=4, + num_moe_experts=4, + moe_grouped_gemm=True, + gated_linear_unit=True, + activation_func=F.silu, + activation_func_clamp_value=10.0, + use_transformer_engine_op_fuser=True, + ) + + assert config.activation_func_clamp_value == 10.0 + assert config.use_transformer_engine_op_fuser is True + + def test_remove_glu_interleaving_restores_contiguous_gate_and_linear_halves(): interleaved = torch.tensor([[1, 2, 5, 6, 3, 4, 7, 8], [11, 12, 15, 16, 13, 14, 17, 18]]) expected = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8], [11, 12, 13, 14, 15, 16, 17, 18]]) @@ -415,11 +432,21 @@ def __init__(self, glu_interleave_size, *, activation_recompute_in_mlp=False): self.activation_recompute_in_mlp = activation_recompute_in_mlp class FakeScaledClampedQGeGLU(torch.nn.Module): - def __init__(self, glu_interleave_size, *, activation_recompute_in_mlp=False, limit=None): + def __init__( + self, + glu_interleave_size, + *, + activation_recompute_in_mlp=False, + limit=None, + alpha=1.702, + glu_linear_offset=1.0, + ): super().__init__() self.glu_interleave_size = glu_interleave_size self.activation_recompute_in_mlp = activation_recompute_in_mlp self.limit = limit + self.alpha = alpha + self.glu_linear_offset = glu_linear_offset class FakeScaledSReLU(torch.nn.Module): def __init__(self, *, activation_recompute_in_mlp=False): @@ -447,8 +474,15 @@ def register_forward_pre_hook(self, hook): ) -def test_make_fused_ops_uses_clamped_qgeglu_for_quick_gelu(monkeypatch): - """quick_gelu + clamp value → ScaledClampedQGeGLU(limit=clamp).""" +@pytest.mark.parametrize( + ("activation_func", "expected_alpha", "expected_offset"), + [(quick_gelu, 1.702, 1.0), (F.silu, 1.0, 0.0)], + ids=("quick-geglu", "clamped-swiglu"), +) +def test_make_fused_ops_uses_clamped_qgeglu( + monkeypatch, activation_func, expected_alpha, expected_offset +): + """Clamped quick GeGLU and SwiGLU use the appropriate TE parameters.""" fake_te, FakeGroupedLinear = _make_fake_te_namespace() monkeypatch.setattr(experts_module, "te", fake_te) @@ -458,10 +492,10 @@ def test_make_fused_ops_uses_clamped_qgeglu_for_quick_gelu(monkeypatch): moe_mlp_glu_interleave_size=4, delay_wgrad_compute=False, activation_func_clamp_value=7.0, - activation_func=quick_gelu, + activation_func=activation_func, gated_linear_unit=True, ) - module.activation_func = quick_gelu + module.activation_func = activation_func module.activation_recompute = True common = dict( device="cuda", @@ -483,6 +517,8 @@ def test_make_fused_ops_uses_clamped_qgeglu_for_quick_gelu(monkeypatch): assert activation.glu_interleave_size == 4 assert activation.activation_recompute_in_mlp is True assert activation.limit == 7.0 + assert activation.alpha == expected_alpha + assert activation.glu_linear_offset == expected_offset def test_make_fused_ops_uses_scaled_srelu_for_weighted_squared_relu(monkeypatch): @@ -573,12 +609,18 @@ def _install_fake_te_ops_modules( def _make_fused_impl_support_module( - FakeGroupedLinear, *, activation_func, gated_linear_unit, use_fused_weighted_squared_relu=False + FakeGroupedLinear, + *, + activation_func, + gated_linear_unit, + use_fused_weighted_squared_relu=False, + activation_func_clamp_value=None, ): module = TEGroupedMLP.__new__(TEGroupedMLP) torch.nn.Module.__init__(module) module.config = SimpleNamespace( activation_func=activation_func, + activation_func_clamp_value=activation_func_clamp_value, gated_linear_unit=gated_linear_unit, use_fused_weighted_squared_relu=use_fused_weighted_squared_relu, moe_apply_probs_on_input=False, @@ -634,6 +676,33 @@ def __init__(self, *args, **kwargs): assert module._is_fused_impl_supported() is False +@pytest.mark.parametrize( + ("te_217_or_later", "include_clamped_qgeglu", "expected"), + [(True, True, True), (False, True, False), (True, False, False)], +) +def test_is_fused_impl_supported_gates_clamped_swiglu( + monkeypatch, te_217_or_later, include_clamped_qgeglu, expected +): + fake_te, FakeGroupedLinear = _make_fake_te_namespace() + monkeypatch.setattr(experts_module, "te", fake_te) + monkeypatch.setattr(experts_module, "HAVE_TE", True) + monkeypatch.setattr( + experts_module, "is_te_min_version", lambda version: version == "2.14.0" or te_217_or_later + ) + _install_fake_te_ops_modules( + monkeypatch, fake_te, include_clamped_qgeglu=include_clamped_qgeglu + ) + + module = _make_fused_impl_support_module( + FakeGroupedLinear, + activation_func=F.silu, + gated_linear_unit=True, + activation_func_clamp_value=10.0, + ) + + assert module._is_fused_impl_supported() is expected + + @pytest.mark.parametrize( ("use_fused_weighted_squared_relu", "gated_linear_unit", "expected"), [(True, False, True), (False, False, False), (True, True, False)], diff --git a/tests/unit_tests/transformer/moe/test_moe_layer.py b/tests/unit_tests/transformer/moe/test_moe_layer.py index c4a3283ac5b..cd7529e6dfb 100644 --- a/tests/unit_tests/transformer/moe/test_moe_layer.py +++ b/tests/unit_tests/transformer/moe/test_moe_layer.py @@ -367,17 +367,19 @@ def test_moe_layer_recompute_forward_backward( requires_grad=True, ) - # Create padding mask if needed: shape [batch_size, sequence_length] + # Sequence-parallel hidden states are TP-local, while the mask arrives + # with the full sequence dimension and is scattered inside MoELayer. padding_mask = None if with_padding_mask: - padding_mask = torch.ones( + mask_sequence_length = sequence_length * tp_size if tp_size > 1 else sequence_length + padding_mask = torch.zeros( micro_batch_size, - sequence_length, + mask_sequence_length, device=torch.cuda.current_device(), dtype=torch.bool, ) # Mark last 4 tokens as padding for each batch - padding_mask[:, -4:] = False + padding_mask[:, -4:] = True output, _ = moe_layer(hidden_states, padding_mask=padding_mask) @@ -399,5 +401,16 @@ def test_moe_layer_recompute_forward_backward( Utils.destroy_model_parallel() + def test_sequence_parallel_padding_mask_alignment(self): + """Exercise global-mask alignment with TP-local sequence-parallel activations.""" + self.test_moe_layer_recompute_forward_backward( + num_moe_experts=2, + moe_token_dispatcher_type="alltoall", + with_padding_mask=True, + tp_size=2, + ep_size=1, + fp8=False, + ) + def teardown_method(self, method): Utils.destroy_model_parallel() diff --git a/tests/unit_tests/transformer/moe/test_routers.py b/tests/unit_tests/transformer/moe/test_routers.py index b215b59cfa0..20a67a7e36f 100644 --- a/tests/unit_tests/transformer/moe/test_routers.py +++ b/tests/unit_tests/transformer/moe/test_routers.py @@ -9,13 +9,15 @@ from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_submodules from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules from megatron.core.transformer.moe.moe_utils import ( + get_default_pg_collection, get_updated_expert_bias, router_gating_linear, topk_routing_with_score_function, ) -from megatron.core.transformer.moe.router import Router +from megatron.core.transformer.moe.router import InferenceTopKRouter, Router, TopKRouter from megatron.core.transformer.spec_utils import get_submodules from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.transformer.transformer_layer import TransformerLayer from megatron.training.initialize import _set_random_seed from tests.unit_tests.test_utilities import Utils @@ -134,9 +136,15 @@ def test_aux_loss(self): @pytest.mark.internal @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_router_with_padding_mask(self): - """Test that padding mask correctly excludes padding tokens from routing.""" + @pytest.mark.parametrize("router_fusion", [False, True]) + def test_router_with_padding_mask(self, router_fusion): + """Test that HybridEP excludes padding tokens from routing.""" + if router_fusion and not HAVE_ROUTER_FUSION: + pytest.skip("TE fused router ops not available") self.router = self.router.cuda() + self.router.config.moe_router_fusion = router_fusion + self.router.config.moe_token_dispatcher_type = "flex" + self.router.config.moe_flex_dispatcher_backend = "hybridep" seq_len = 32 batch_size = 2 hidden_size = self.router.config.hidden_size @@ -176,9 +184,50 @@ def test_router_with_padding_mask(self): self.router.config.num_moe_experts, ) + padding_rows = padding_mask.reshape(-1) + assert torch.count_nonzero(probs_with_mask[padding_rows]) == 0 + assert not routing_map_with_mask[padding_rows].any() + # Verify that probs for valid tokens are similar assert torch.equal(probs_valid_part, probs_without_mask) + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize( + "dispatcher,backend,capacity_factor,rank_capacity_factor", + [ + ("allgather", "deepep", None, None), + ("alltoall", "deepep", None, None), + ("flex", "deepep", None, None), + ("flex", "deepepv2", None, None), + ("flex", "hybridep", 1.0, None), + ("flex", "hybridep", None, 1.0), + ], + ) + def test_padding_mask_preserves_routes_outside_dropless_hybridep( + self, dispatcher, backend, capacity_factor, rank_capacity_factor + ): + """Only dropless HybridEP may consume a sparse route map.""" + self.router = self.router.cuda() + self.router.config.moe_token_dispatcher_type = dispatcher + self.router.config.moe_flex_dispatcher_backend = backend + self.router.config.moe_expert_capacity_factor = capacity_factor + self.router.config.moe_expert_rank_capacity_factor = rank_capacity_factor + hidden_states = torch.randn( + (16, 2, self.router.config.hidden_size), device="cuda", dtype=torch.bfloat16 + ) + padding_mask = torch.zeros((16, 2), dtype=torch.bool, device="cuda") + padding_mask[8:, :] = True + + with torch.no_grad(): + probs_with_mask, routing_map_with_mask = self.router( + hidden_states, padding_mask=padding_mask + ) + probs_without_mask, routing_map_without_mask = self.router(hidden_states) + + torch.testing.assert_close(probs_with_mask, probs_without_mask) + assert torch.equal(routing_map_with_mask, routing_map_without_mask) + @pytest.mark.internal @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_router_dtype(self): @@ -463,6 +512,30 @@ def test_router_forward_aux_free(self): # Print some debug info print("Updated bias after first forward pass:", updated_bias) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_expert_bias_ignores_2d_padding_mask(self): + self.router = self.router.cuda() + routing_map = torch.tensor( + [ + [True, False, False, False, False, False, False, False], + [False, True, False, False, False, False, False, False], + [False, False, True, False, False, False, False, False], + [False, False, False, True, False, False, False, False], + [False, False, False, False, True, False, False, False], + [False, False, False, False, False, True, False, False], + ], + device="cuda", + ) + padding_mask = torch.tensor( + [[False, True, False], [True, False, True]], device="cuda" + ) + self.router.local_tokens_per_expert.zero_() + + self.router._apply_expert_bias(routing_map, padding_mask) + + expected = routing_map[~padding_mask.reshape(-1)].sum(dim=0) + assert torch.equal(self.router.local_tokens_per_expert, expected) + @pytest.mark.internal @pytest.mark.skipif( not torch.cuda.is_available() or not HAVE_ROUTER_FUSION, @@ -615,3 +688,151 @@ def test_topk_routing_precomputed_indices_equivalence(score_function, use_pre_so ) expected_map = torch.zeros_like(logits, dtype=torch.bool).scatter(1, alt_indices, True) assert torch.equal(map_alt, expected_map) + + +def _hash_routing_config(**overrides): + """Create a small TransformerConfig for hash-routing tests.""" + defaults = dict( + num_layers=3, + hidden_size=16, + num_attention_heads=8, + num_moe_experts=4, + moe_router_topk=2, + moe_router_load_balancing_type="aux_loss", + moe_aux_loss_coeff=0.0, + moe_router_dtype="fp32", + add_bias_linear=False, + use_cpu_initialization=True, + moe_n_hash_layers=2, + actual_vocab_size=128, + ) + defaults.update(overrides) + return TransformerConfig(**defaults) + + +def _make_hash_router(config, layer_number, *, inference=False, is_mtp_layer=False): + router_cls = InferenceTopKRouter if inference else TopKRouter + router = router_cls( + config=config, pg_collection=get_default_pg_collection(), is_mtp_layer=is_mtp_layer + ) + router.set_layer_number(layer_number) + return router + + +class TestHashRouting: + """Hash expert selection and generic TransformerLayer integration.""" + + def setup_method(self, method): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + expert_model_parallel_size=1, + ) + _set_random_seed(seed_=42, data_parallel_random_init=False) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) + def test_hash_routing_correctness(self, score_function): + config = _hash_routing_config(moe_router_score_function=score_function) + router = _make_hash_router(config, layer_number=1).cuda() + + logits = torch.randn(16, config.num_moe_experts, device="cuda") + input_ids = torch.randint(0, config.actual_vocab_size, (4, 4), device="cuda") + routing_probs, routing_map = router._hash_routing(logits, input_ids) + + if score_function == "softmax": + scores = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits) + elif score_function == "sigmoid": + scores = torch.sigmoid(logits.float()).type_as(logits) + else: + scores = torch.nn.functional.softplus(logits.float()).sqrt().type_as(logits) + + top_indices = router.tid2eid[input_ids.T.reshape(-1)].long() + expected_probs = scores.gather(1, top_indices) + if score_function != "softmax": + expected_probs = expected_probs / (expected_probs.sum(dim=-1, keepdim=True) + 1e-20) + + expected_map = torch.zeros_like(routing_map).scatter(1, top_indices, True) + expected_routing_probs = torch.zeros_like(routing_probs).scatter( + 1, top_indices, expected_probs + ) + assert torch.equal(routing_map, expected_map) + torch.testing.assert_close(routing_probs, expected_routing_probs) + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize( + ("force_option", "force_value"), + [("moe_router_force_load_balancing", True), ("moe_router_force_biased", 0.5)], + ) + def test_forced_routing_overrides_hash_table(self, force_option, force_value): + config = _hash_routing_config(moe_router_score_function="softmax") + router = _make_hash_router(config, layer_number=1).cuda() + setattr(router.config, force_option, force_value) + + router.tid2eid.fill_(0) + logits = torch.tensor([[0.0, 1.0, 8.0, 9.0]], device="cuda").expand(16, -1) + input_ids = torch.zeros((4, 4), dtype=torch.long, device="cuda") + _, routing_map = router._hash_routing(logits, input_ids) + + forced_indices = torch.topk(logits, k=router.topk, dim=1).indices + expected_map = torch.zeros_like(routing_map).scatter(1, forced_indices, True) + assert torch.equal(routing_map, expected_map) + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_hash_layer_selection(self): + config = _hash_routing_config( + moe_router_enable_expert_bias=True, moe_router_score_function="sigmoid" + ) + first = _make_hash_router(config, layer_number=1) + boundary = _make_hash_router(config, layer_number=2) + learned = _make_hash_router(config, layer_number=3) + mtp = _make_hash_router(config, layer_number=1, is_mtp_layer=True) + + assert first.is_hash_layer and first.tid2eid is not None + assert boundary.is_hash_layer and boundary.tid2eid is not None + assert not learned.is_hash_layer and learned.tid2eid is None + assert not mtp.is_hash_layer and mtp.tid2eid is None + assert "tid2eid" in first.state_dict() + assert "tid2eid" not in learned.state_dict() + assert first.enable_expert_bias is False + assert learned.enable_expert_bias is True + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_inference_router_hash_fallback(self): + config = _hash_routing_config(moe_router_score_function="sqrtsoftplus") + router = _make_hash_router(config, layer_number=1, inference=True).cuda() + hidden_states = torch.randn(4, 4, config.hidden_size, device="cuda") + input_ids = torch.randint(0, config.actual_vocab_size, (4, 4), device="cuda") + + routing_probs, routing_map = router(hidden_states, input_ids=input_ids) + + assert router.is_hash_layer + assert routing_probs.shape == routing_map.shape == (16, config.num_moe_experts) + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("moe_layer_recompute", [False, True]) + def test_transformer_layer_hash_routing(self, moe_layer_recompute): + config = _hash_routing_config(moe_n_hash_layers=1, moe_layer_recompute=moe_layer_recompute) + submodules = get_gpt_layer_local_submodules( + num_experts=config.num_moe_experts, moe_grouped_gemm=False + ) + layer = TransformerLayer(config, submodules, layer_number=1).cuda() + hidden_states = torch.randn(8, 2, config.hidden_size, device="cuda", requires_grad=True) + input_ids = torch.randint(0, config.actual_vocab_size, (2, 8), device="cuda") + + output, _ = layer(hidden_states=hidden_states, attention_mask=None, input_ids=input_ids) + + assert layer.mlp.router.is_hash_layer + assert output.shape == hidden_states.shape + assert torch.isfinite(output).all() + output.sum().backward() + assert hidden_states.grad is not None + assert torch.isfinite(hidden_states.grad).all() diff --git a/tests/unit_tests/transformer/moe/test_shared_experts.py b/tests/unit_tests/transformer/moe/test_shared_experts.py index 8c84aae7097..8fa40800b88 100644 --- a/tests/unit_tests/transformer/moe/test_shared_experts.py +++ b/tests/unit_tests/transformer/moe/test_shared_experts.py @@ -52,6 +52,15 @@ def __init__(self, glu_interleave_size): self.glu_interleave_size = glu_interleave_size +class _FakeTEScaledClampedQGeGLU(torch.nn.Module): + def __init__(self, glu_interleave_size, *, alpha, limit, glu_linear_offset): + super().__init__() + self.glu_interleave_size = glu_interleave_size + self.alpha = alpha + self.limit = limit + self.glu_linear_offset = glu_linear_offset + + class _FakeTESequential(torch.nn.Module): def append(self, module): self.add_module(str(len(self._modules)), module) @@ -90,6 +99,7 @@ def _fake_te_module(linear_cls=_FakeTELinear): ops=SimpleNamespace( GroupedLinear=_FakeTEGroupedLinear, ScaledSwiGLU=_FakeTEScaledSwiGLU, + ScaledClampedQGeGLU=_FakeTEScaledClampedQGeGLU, Sequential=_FakeTESequential, ), fp8_autocast=_FakeFP8Autocast, @@ -123,6 +133,7 @@ def _fake_shared_expert(**config_kwargs): add_bias_linear=False, gated_linear_unit=True, activation_func=F.silu, + activation_func_clamp_value=None, moe_shared_expert_glu_interleave_size=32, delay_wgrad_compute=False, sequence_parallel=False, @@ -183,6 +194,19 @@ def test_validate_fused_grouped_swiglu_requires_te(monkeypatch): shared_expert._validate_fused_grouped_swiglu() +@pytest.mark.parametrize("has_clamped_op", [True, False]) +def test_validate_fused_grouped_swiglu_requires_clamped_te_support(monkeypatch, has_clamped_op): + fake_te = _patch_fake_shared_expert_te(monkeypatch) + if has_clamped_op: + monkeypatch.setattr(shared_experts_module, "is_te_min_version", lambda *_: False) + else: + del fake_te.pytorch.ops.ScaledClampedQGeGLU + shared_expert = _fake_shared_expert(activation_func_clamp_value=7.0) + + with pytest.raises(RuntimeError, match="ScaledClampedQGeGLU"): + shared_expert._validate_fused_grouped_swiglu() + + @pytest.mark.parametrize( ("config_kwargs", "bad_linear", "match"), [ @@ -238,6 +262,21 @@ def test_make_fused_grouped_swiglu_ops_builds_grouped_pipeline(monkeypatch): assert fc2_op.weight0 is shared_expert.linear_fc2.weight +def test_make_fused_grouped_swiglu_ops_builds_clamped_activation(monkeypatch): + _patch_fake_shared_expert_te(monkeypatch) + shared_expert = _fake_shared_expert(activation_func_clamp_value=7.0) + + shared_expert._validate_fused_grouped_swiglu() + ops = shared_expert._make_fused_grouped_swiglu_ops() + + activation_op = list(ops.children())[1] + assert isinstance(activation_op, _FakeTEScaledClampedQGeGLU) + assert activation_op.glu_interleave_size == 32 + assert activation_op.alpha == 1.0 + assert activation_op.limit == 7.0 + assert activation_op.glu_linear_offset == 0.0 + + def test_fused_grouped_swiglu_ops_replay_linear_pre_forward_hooks(monkeypatch): _patch_fake_shared_expert_te(monkeypatch) shared_expert = _fake_shared_expert() @@ -408,3 +447,72 @@ def test_shared_expert_forward_backward(self, dispatcher_type: str, tp_size, ep_ assert torch.allclose( p_overlap.grad, p_no_overlap.grad ), f"max diff: {torch.max(torch.abs(p_overlap.grad - p_no_overlap.grad))}" + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("bias_activation_fusion", [True, False]) + def test_shared_expert_clamped_swiglu(self, bias_activation_fusion): + """Verify clamped SwiGLU parity for overlapped and synchronous shared experts.""" + Utils.initialize_model_parallel(tensor_model_parallel_size=1, expert_model_parallel_size=1) + clamp_value = 1.0 + + _set_random_seed(seed_=123, data_parallel_random_init=False) + moe_layer_overlap = self.get_moe_layer( + moe_shared_expert_overlap=True, + moe_token_dispatcher_type="alltoall", + activation_func_clamp_value=clamp_value, + bias_activation_fusion=bias_activation_fusion, + ).to(dtype=torch.bfloat16) + + _set_random_seed(seed_=123, data_parallel_random_init=False) + moe_layer_no_overlap = self.get_moe_layer( + moe_shared_expert_overlap=False, + moe_token_dispatcher_type="alltoall", + activation_func_clamp_value=clamp_value, + bias_activation_fusion=bias_activation_fusion, + ).to(dtype=torch.bfloat16) + moe_layer_no_overlap.load_state_dict(moe_layer_overlap.state_dict()) + + hidden_states = ( + torch.randn((32, 2, self.config.hidden_size), device="cuda", dtype=torch.bfloat16) * 5.0 + ).requires_grad_(True) + hidden_states_no_overlap = hidden_states.detach().clone().requires_grad_(True) + + output_overlap, _ = moe_layer_overlap(hidden_states) + output_no_overlap, _ = moe_layer_no_overlap(hidden_states_no_overlap) + + cos_out = F.cosine_similarity( + output_overlap.flatten().unsqueeze(0).float(), + output_no_overlap.flatten().unsqueeze(0).float(), + ).item() + assert cos_out > 0.999, ( + f"shared-expert clamp output mismatch (fusion={bias_activation_fusion}): " + f"cos sim = {cos_out:.6f}" + ) + + output_overlap.mean().backward() + output_no_overlap.mean().backward() + + for p_overlap, p_no_overlap in zip( + moe_layer_overlap.parameters(), moe_layer_no_overlap.parameters() + ): + assert torch.allclose(p_overlap.grad, p_no_overlap.grad), ( + f"shared-expert clamp mismatch (fusion={bias_activation_fusion}); " + f"max diff: {torch.max(torch.abs(p_overlap.grad - p_no_overlap.grad))}" + ) + + _set_random_seed(seed_=123, data_parallel_random_init=False) + moe_layer_unclamped = self.get_moe_layer( + moe_shared_expert_overlap=False, + moe_token_dispatcher_type="alltoall", + activation_func_clamp_value=None, + bias_activation_fusion=bias_activation_fusion, + ).to(dtype=torch.bfloat16) + moe_layer_unclamped.load_state_dict(moe_layer_overlap.state_dict()) + + hidden_states_unclamped = hidden_states.detach().clone().requires_grad_(True) + output_unclamped, _ = moe_layer_unclamped(hidden_states_unclamped) + assert not torch.allclose(output_no_overlap, output_unclamped), ( + "Clamping had no observable effect on shared-expert output; " + "activation_func_clamp_value may not be plumbed through." + ) diff --git a/tests/unit_tests/transformer/moe/test_token_dispatcher.py b/tests/unit_tests/transformer/moe/test_token_dispatcher.py index 20839310ce2..cb06e7deb11 100644 --- a/tests/unit_tests/transformer/moe/test_token_dispatcher.py +++ b/tests/unit_tests/transformer/moe/test_token_dispatcher.py @@ -7,8 +7,10 @@ from megatron.core import config, parallel_state from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_submodules +from megatron.core.transformer.moe.fused_a2a import HYBRIDEP_TOKEN_ALIGNMENT, reset_hybrid_ep_buffer from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules from megatron.core.transformer.moe.moe_utils import get_capacity +from megatron.core.transformer.moe.token_dispatcher import _HybridEPManager from megatron.core.transformer.spec_utils import get_submodules from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.typed_torch import apply_module @@ -426,6 +428,98 @@ def is_nccl_ep_available(): return HAVE_TE_EP +def test_hybridep_pad_uneven_dispatch_inputs_metadata(monkeypatch): + manager = _HybridEPManager.__new__(_HybridEPManager) + manager.group = object() + manager.num_local_experts = 2 + manager.num_experts = 4 + manager.config = TransformerConfig( + num_layers=1, + hidden_size=16, + num_attention_heads=4, + num_moe_experts=4, + moe_router_topk=2, + moe_hybridep_pad_uneven_dispatch_inputs=True, + ) + manager.moe_expert_rank_capacity_factor = None + manager.drop_and_pad = False + + local_num_tokens = 17 + max_num_tokens_across_ep = 70 + padded_num_tokens = ( + max_num_tokens_across_ep + -max_num_tokens_across_ep % HYBRIDEP_TOKEN_ALIGNMENT + ) + routing_map = torch.ones((local_num_tokens, manager.num_experts), dtype=torch.bool) + probs = torch.ones((local_num_tokens, manager.num_experts), dtype=torch.float32) + + def fake_all_reduce(tensor, op=None, group=None): + assert op == torch.distributed.ReduceOp.MAX + assert group is manager.group + tensor.fill_(max_num_tokens_across_ep) + + monkeypatch.setattr(torch.distributed, "all_reduce", fake_all_reduce) + + manager.setup_metadata(routing_map, probs) + + assert manager._original_num_tokens == local_num_tokens + assert manager._padded_num_tokens == padded_num_tokens + assert manager.routing_map.shape == (padded_num_tokens, manager.num_experts) + assert manager.token_probs.shape == (padded_num_tokens, manager.num_experts) + torch.testing.assert_close(manager.routing_map[:local_num_tokens], routing_map) + torch.testing.assert_close(manager.token_probs[:local_num_tokens], probs) + assert not manager.routing_map[local_num_tokens:].any() + assert not manager.token_probs[local_num_tokens:].any() + + +@pytest.mark.parametrize( + "is_capturing,is_compiling", + [(True, False), (False, True)], +) +def test_hybridep_static_sequence_packing_metadata_skips_collective( + monkeypatch, is_capturing, is_compiling +): + manager = _HybridEPManager.__new__(_HybridEPManager) + manager.group = object() + manager.num_local_experts = 2 + manager.num_experts = 4 + manager.config = TransformerConfig( + num_layers=1, + hidden_size=16, + num_attention_heads=4, + num_moe_experts=4, + moe_router_topk=2, + moe_hybridep_pad_uneven_dispatch_inputs=True, + ) + # Static sequence packing is normally validated together with TE and CUDA + # graph configuration. This unit test isolates HybridEP metadata behavior. + manager.config.sequence_packing_scheduler = "dp_balanced" + manager.moe_expert_rank_capacity_factor = None + manager.drop_and_pad = False + + local_num_tokens = 17 + padded_num_tokens = local_num_tokens + -local_num_tokens % HYBRIDEP_TOKEN_ALIGNMENT + routing_map = torch.ones((local_num_tokens, manager.num_experts), dtype=torch.bool) + probs = torch.ones((local_num_tokens, manager.num_experts), dtype=torch.float32) + + def fail_all_reduce(*args, **kwargs): + pytest.fail("static capture/compile metadata must not issue all_reduce") + + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: is_capturing) + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: is_compiling) + monkeypatch.setattr(torch.distributed, "all_reduce", fail_all_reduce) + + manager.setup_metadata(routing_map, probs) + + assert manager._original_num_tokens == local_num_tokens + assert manager._padded_num_tokens == padded_num_tokens + assert manager.routing_map.shape == (padded_num_tokens, manager.num_experts) + assert manager.token_probs.shape == (padded_num_tokens, manager.num_experts) + torch.testing.assert_close(manager.routing_map[:local_num_tokens], routing_map) + torch.testing.assert_close(manager.token_probs[:local_num_tokens], probs) + assert not manager.routing_map[local_num_tokens:].any() + assert not manager.token_probs[local_num_tokens:].any() + + @pytest.mark.skipif( not is_deep_ep_available() and not is_hybrid_ep_available(), reason="Deep EP and Hybrid EP are not available", @@ -435,6 +529,7 @@ def setup_method(self, method): pass def teardown_method(self, method): + reset_hybrid_ep_buffer() Utils.destroy_model_parallel() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index ef556b29f42..2c31588f4aa 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -16,7 +16,7 @@ get_gpt_mtp_block_spec, ) from megatron.core.models.gpt.gpt_model import GPTModel -from megatron.core.models.hybrid.hybrid_block import HybridStack +from megatron.core.models.hybrid.hybrid_block import HybridStack, HyperConnectionHybridLayer from megatron.core.models.hybrid.hybrid_layer_allocation import validate_segment_layers from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec from megatron.core.num_microbatches_calculator import ( @@ -35,11 +35,12 @@ CudaGraphManager, TECudaGraphHelper, _CudagraphGlobalRecord, + _layer_is_graphable, create_cudagraphs, ) from megatron.core.transformer.enums import CudaGraphModule, CudaGraphScope, InferenceCudaGraphScope from megatron.core.transformer.mlp import MLPSubmodules -from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.module import GraphableMegatronModule, MegatronModule from megatron.core.transformer.moe.fused_a2a import reset_hybrid_ep_buffer from megatron.core.transformer.spec_utils import ModuleSpec, get_submodules from megatron.core.transformer.transformer_block import TransformerBlock @@ -973,6 +974,45 @@ def test_gpu_cudagraph(self): del parallel_mamba_block.layers[_].cudagraph_manager.cudagraph_runners[0].fwd_graph +class TestMhcHybridCudagraphDiscovery: + def setup_method(self, method): + initialize_rng_tracker(use_te_rng_tracker=True, force_reset=True) + Utils.initialize_model_parallel(tensor_model_parallel_size=2) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def test_wrapped_hybrid_layers_are_graphable(self): + assert issubclass(HyperConnectionHybridLayer, GraphableMegatronModule) + + layer_type_list = validate_segment_layers("*-") + config = TransformerConfig( + hidden_size=256, + num_layers=len(layer_type_list), + num_attention_heads=4, + use_cpu_initialization=True, + cuda_graph_impl="transformer_engine", + enable_hyper_connections=True, + num_residual_streams=4, + cuda_graph_modules=[CudaGraphModule.attn, CudaGraphModule.mlp], + ) + block = HybridStack( + config, + hybrid_stack_spec.submodules, + layer_type_list=layer_type_list, + pp_layer_offset=0, + pg_collection=ProcessGroupCollection.use_mpu_process_groups( + required_pgs=["tp", "pp", "cp"] + ), + ) + + assert all( + isinstance(layer, HyperConnectionHybridLayer) for layer in block.layers + ) + assert any(_layer_is_graphable(layer, config) for layer in block.layers) + + # Global storage for comparing unique buffer counts across different num_microbatches, # keyed by (pp_size, vpp_size) _unique_buffer_counts = {} diff --git a/tests/unit_tests/transformer/test_hyper_connection_recompute.py b/tests/unit_tests/transformer/test_hyper_connection_recompute.py new file mode 100644 index 00000000000..d9ed2cfa7e0 --- /dev/null +++ b/tests/unit_tests/transformer/test_hyper_connection_recompute.py @@ -0,0 +1,434 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +""" +Unit tests for HyperConnection block-level recomputation. + +Tests the following functionality: +1. HyperConnectionModule._forward_with_checkpoint correctness +2. HyperConnectionModule.apply_h_post with CheckpointManager +3. Multiple HyperConnectionModules chained with a single CheckpointManager +4. Partial checkpoint (last layer not checkpointed) +5. TransformerConfig 'mhc' in recompute_modules option +""" + +import pytest +import torch + +from megatron.core.tensor_parallel.random import CheckpointManager, model_parallel_cuda_manual_seed +from megatron.core.transformer.hyper_connection import HyperConnectionModule +from megatron.core.transformer.transformer_config import TransformerConfig +from tests.unit_tests.test_utilities import Utils + + +class TestHyperConnectionCheckpoint: + """Test HyperConnectionModule checkpoint functionality.""" + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def _create_hyper_connection_module(self, hidden_size=64, num_residual_streams=4): + """Create a HyperConnectionModule for testing.""" + config = TransformerConfig( + num_layers=2, + hidden_size=hidden_size, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=num_residual_streams, + mhc_sinkhorn_iterations=5, # Fewer iterations for faster tests + mhc_init_gating_factor=0.01, + ) + module = HyperConnectionModule(config=config, layer_number=1) + module.cuda() + return module + + def test_apply_h_res_uses_h_res_transpose(self): + """apply_h_res should compute H_res.T @ residual.""" + module = self._create_hyper_connection_module(hidden_size=4, num_residual_streams=2) + h_res = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]], device='cuda') + residual = torch.tensor([[[10.0, 100.0, 3.0, 4.0, 1.0, 2.0, 5.0, 6.0]]], device='cuda') + expected = torch.tensor( + [[[13.0, 106.0, 18.0, 22.0, 24.0, 208.0, 26.0, 32.0]]], device='cuda' + ) + + mixed = module.apply_h_res(h_res, residual) + + torch.testing.assert_close(mixed, expected, atol=0.0, rtol=0.0) + + def test_forward_preserves_three_tuple_api_and_hybrid_can_request_residual(self): + module = self._create_hyper_connection_module(hidden_size=8, num_residual_streams=2) + hidden_states = torch.randn(4, 1, 16, device='cuda', requires_grad=True) + + compatible_output = module(hidden_states) + hybrid_output = module(hidden_states, return_residual=True) + + assert len(compatible_output) == 3 + assert len(hybrid_output) == 4 + for compatible, hybrid in zip(compatible_output, hybrid_output[:3]): + torch.testing.assert_close(compatible, hybrid) + assert hybrid_output[3].shape == hidden_states.shape + + def test_forward_normal_vs_checkpoint_correctness(self): + """ + Test that _forward_with_checkpoint produces the same outputs as _forward_normal. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + module = self._create_hyper_connection_module(hidden_size, num_streams) + + # Create input tensors + hidden_states = torch.randn( + seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True + ) + residual = torch.randn( + seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True + ) + + # Clone inputs for comparison + hidden_states_ckpt = hidden_states.detach().clone().requires_grad_(True) + residual_ckpt = residual.detach().clone().requires_grad_(True) + + # Forward without checkpoint (reference) + torch.manual_seed(42) + torch.cuda.manual_seed(42) + aggregated_ref, h_res_ref, h_post_ref, residual_ref = module._forward_normal(hidden_states) + mixed_ref = module.apply_h_res(h_res_ref, residual) + loss_ref = aggregated_ref.sum() + mixed_ref.sum() + h_post_ref.sum() + loss_ref.backward() + grad_hidden_ref = hidden_states.grad.clone() + grad_residual_ref = residual.grad.clone() + + # Forward with checkpoint + torch.manual_seed(42) + torch.cuda.manual_seed(42) + manager = CheckpointManager() + aggregated_ckpt, h_res_ckpt, h_post_ckpt, residual_ckpt_out = ( + module._forward_with_checkpoint(hidden_states_ckpt, manager) + ) + mixed_ckpt = module.apply_h_res(h_res_ckpt, residual_ckpt) + # Calculate loss before discarding outputs + loss_ckpt = aggregated_ckpt.sum() + mixed_ckpt.sum() + h_post_ckpt.sum() + + # Register unified recompute hook + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + # Backward pass + loss_ckpt.backward() + grad_hidden_ckpt = hidden_states_ckpt.grad.clone() + grad_residual_ckpt = residual_ckpt.grad.clone() + + # Verify gradients match + assert torch.allclose(grad_hidden_ckpt, grad_hidden_ref, atol=1e-5), ( + f"Hidden states gradients mismatch:\n" + f"Checkpoint: {grad_hidden_ckpt}\n" + f"Reference: {grad_hidden_ref}" + ) + assert torch.allclose(grad_residual_ckpt, grad_residual_ref, atol=1e-5), ( + f"Residual gradients mismatch:\n" + f"Checkpoint: {grad_residual_ckpt}\n" + f"Reference: {grad_residual_ref}" + ) + + def test_apply_h_post_with_checkpoint(self): + """ + Test that apply_h_post with manager produces correct gradients. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + module = self._create_hyper_connection_module(hidden_size, num_streams) + + # Create input tensors + x = torch.randn(seq_len, batch_size, hidden_size, device='cuda', requires_grad=True) + bias = torch.randn(hidden_size, device='cuda') + h_post = torch.randn(seq_len, batch_size, num_streams, device='cuda', requires_grad=True) + + # Clone inputs + x_ckpt = x.detach().clone().requires_grad_(True) + h_post_ckpt = h_post.detach().clone().requires_grad_(True) + + # Reference: without checkpoint (manager=None) + torch.manual_seed(42) + x_out_ref, bias_out_ref = module.apply_h_post((x, bias), h_post, manager=None) + loss_ref = x_out_ref.sum() + if bias_out_ref is not None: + loss_ref = loss_ref + bias_out_ref.sum() + loss_ref.backward() + grad_x_ref = x.grad.clone() + grad_h_post_ref = h_post.grad.clone() + + # With checkpoint (manager provided) + torch.manual_seed(42) + manager = CheckpointManager() + x_out_ckpt, bias_out_ckpt = module.apply_h_post( + (x_ckpt, bias), h_post_ckpt, manager=manager + ) + loss_ckpt = x_out_ckpt.sum() + if bias_out_ckpt is not None: + loss_ckpt = loss_ckpt + bias_out_ckpt.sum() + + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + loss_ckpt.backward() + grad_x_ckpt = x_ckpt.grad.clone() + grad_h_post_ckpt = h_post_ckpt.grad.clone() + + # Verify gradients + assert torch.allclose(grad_x_ckpt, grad_x_ref, atol=1e-5) + assert torch.allclose(grad_h_post_ckpt, grad_h_post_ref, atol=1e-5) + + def test_forward_with_manager_parameter(self): + """ + Test forward() method with mhc_recompute_manager parameter. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + module = self._create_hyper_connection_module(hidden_size, num_streams) + + # Create input tensors + hidden_states = torch.randn( + seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True + ) + + # Clone inputs + hidden_states_ckpt = hidden_states.detach().clone().requires_grad_(True) + + # Reference: forward without manager (uses _forward_normal) + torch.manual_seed(42) + torch.cuda.manual_seed(42) + aggregated_ref, h_res_ref, h_post_ref = module.forward( + hidden_states, mhc_recompute_manager=None + ) + loss_ref = aggregated_ref.sum() + h_res_ref.sum() + h_post_ref.sum() + loss_ref.backward() + grad_hidden_ref = hidden_states.grad.clone() + + # With manager (uses _forward_with_checkpoint) + torch.manual_seed(42) + torch.cuda.manual_seed(42) + manager = CheckpointManager() + aggregated_ckpt, h_res_ckpt, h_post_ckpt = module.forward( + hidden_states_ckpt, mhc_recompute_manager=manager + ) + loss_ckpt = aggregated_ckpt.sum() + h_res_ckpt.sum() + h_post_ckpt.sum() + + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + loss_ckpt.backward() + grad_hidden_ckpt = hidden_states_ckpt.grad.clone() + + # Verify gradients match + assert torch.allclose(grad_hidden_ckpt, grad_hidden_ref, atol=1e-5) + + +class TestMHCBlockRecomputeIntegration: + """Test CheckpointManager integration with HyperConnection.""" + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def test_multiple_hyper_connections_in_chain(self): + """ + Test that multiple HyperConnectionModules can be chained together + with a single CheckpointManager. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + n_channels = num_streams * hidden_size + + # Create multiple HyperConnection modules (simulating multiple layers) + config = TransformerConfig( + num_layers=4, + hidden_size=hidden_size, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=num_streams, + mhc_sinkhorn_iterations=5, + mhc_init_gating_factor=0.01, + ) + + modules = [ + HyperConnectionModule(config=config, layer_number=i + 1).cuda() for i in range(3) + ] + + # Create input tensors + hidden_states_ref = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + residual_ref = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + + hidden_states_ckpt = hidden_states_ref.detach().clone().requires_grad_(True) + residual_ckpt = residual_ref.detach().clone().requires_grad_(True) + + # Reference: forward without checkpoint + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + h = hidden_states_ref + r = residual_ref + for module in modules: + agg, h_res, h_post = module.forward(h, mhc_recompute_manager=None) + agg, _ = module.apply_h_post((0.1 * agg, None), h_post, manager=None) + mixed = module.apply_h_res(h_res, r) # Apply h_res to get mixed [s, b, n*C] + h = agg + mixed + r = h + + loss_ref = h.sum() + loss_ref.backward() + grad_hidden_ref = hidden_states_ref.grad.clone() + grad_residual_ref = residual_ref.grad.clone() + + # With checkpoint using single manager + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + manager = CheckpointManager() + + h = hidden_states_ckpt + r = residual_ckpt + for module in modules: + agg, h_res, h_post = module.forward(h, mhc_recompute_manager=manager) + agg, _ = module.apply_h_post((0.1 * agg, None), h_post, manager=manager) + mixed = module.apply_h_res(h_res, r) # Apply h_res to get mixed [s, b, n*C] + h = agg + mixed + r = h + + loss_ckpt = h.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + loss_ckpt.backward() + + grad_hidden_ckpt = hidden_states_ckpt.grad.clone() + grad_residual_ckpt = residual_ckpt.grad.clone() + + # Verify gradients + assert torch.allclose( + grad_hidden_ckpt, grad_hidden_ref, atol=1e-4 + ), f"Chained HyperConnection hidden gradients mismatch" + assert torch.allclose( + grad_residual_ckpt, grad_residual_ref, atol=1e-4 + ), f"Chained HyperConnection residual gradients mismatch" + + def test_partial_checkpoint_last_layer_not_checkpointed(self): + """ + Test that when is_last_layer_in_block=True, the final output is NOT checkpointed. + This simulates the TransformerBlock behavior where the last layer's MLP BDA + serves as the hook_tensor for unified recompute. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + config = TransformerConfig( + num_layers=2, + hidden_size=hidden_size, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=num_streams, + mhc_sinkhorn_iterations=5, + mhc_init_gating_factor=0.01, + ) + + module = HyperConnectionModule(config=config, layer_number=1).cuda() + + hidden_states_ref = torch.randn( + seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True + ) + residual_ref = torch.randn( + seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True + ) + + hidden_states_ckpt = hidden_states_ref.detach().clone().requires_grad_(True) + residual_ckpt = residual_ref.detach().clone().requires_grad_(True) + + # Reference + torch.manual_seed(42) + torch.cuda.manual_seed(42) + aggregated_ref, h_res_ref, h_post_ref = module.forward( + hidden_states_ref, mhc_recompute_manager=None + ) + aggregated_ref, _ = module.apply_h_post( + (0.1 * aggregated_ref, None), h_post_ref, manager=None + ) + mixed_ref = module.apply_h_res( + h_res_ref, residual_ref + ) # Apply h_res to get mixed [s, b, n*C] + # Simulate BDA that is NOT checkpointed (last layer) + output_ref = aggregated_ref + 0.5 * mixed_ref + loss_ref = output_ref.sum() + loss_ref.backward() + grad_hidden_ref = hidden_states_ref.grad.clone() + + # With manager - checkpoint everything except final output + torch.manual_seed(42) + torch.cuda.manual_seed(42) + manager = CheckpointManager() + aggregated_ckpt, h_res_ckpt, h_post_ckpt = module.forward( + hidden_states_ckpt, mhc_recompute_manager=manager + ) + + aggregated_ckpt, _ = module.apply_h_post( + (0.1 * aggregated_ckpt, None), h_post_ckpt, manager=manager + ) + mixed_ckpt = module.apply_h_res( + h_res_ckpt, residual_ckpt + ) # Apply h_res to get mixed [s, b, n*C] + # Simulate BDA that is NOT checkpointed (last layer) - this is the hook_tensor + output_ckpt = aggregated_ckpt + 0.5 * mixed_ckpt + + # Register unified recompute on the output (which is not checkpointed) + manager.discard_all_outputs_and_register_unified_recompute(output_ckpt) + + loss_ckpt = output_ckpt.sum() + loss_ckpt.backward() + grad_hidden_ckpt = hidden_states_ckpt.grad.clone() + + # Verify gradients match + assert torch.allclose(grad_hidden_ckpt, grad_hidden_ref, atol=1e-5) + + +class TestTransformerConfigRecomputeMhc: + """Test 'mhc' in recompute_modules configuration.""" + + def test_config_default_value(self): + """Test that 'mhc' is not in recompute_modules by default.""" + config = TransformerConfig(num_layers=2, hidden_size=64, num_attention_heads=4) + assert "mhc" not in config.recompute_modules + + def test_config_enable_mhc_recompute(self): + """Test enabling 'mhc' in recompute_modules.""" + config = TransformerConfig( + num_layers=2, + hidden_size=64, + num_attention_heads=4, + enable_hyper_connections=True, + num_residual_streams=4, + recompute_modules=["core_attn", "mhc"], + recompute_granularity='selective', + ) + assert "mhc" in config.recompute_modules + assert config.enable_hyper_connections is True + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/transformer/test_mhc_block_manager.py b/tests/unit_tests/transformer/test_mhc_block_manager.py new file mode 100644 index 00000000000..0d4f40bba7d --- /dev/null +++ b/tests/unit_tests/transformer/test_mhc_block_manager.py @@ -0,0 +1,522 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import pytest +import torch + +from megatron.core.tensor_parallel.random import ( + CheckpointManager, + CheckpointWithoutOutput, + CheckpointWithoutOutputManager, + initialize_rng_tracker, +) +from tests.unit_tests.test_utilities import Utils + + +class TestCheckpointWithoutOutputManagerAPI: + """Test CheckpointWithoutOutput integration with CheckpointWithoutOutputManager.""" + + def setup_method(self, method): + Utils.initialize_model_parallel() + initialize_rng_tracker(force_reset=True) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def test_reviewed_manager_name_is_compatible_alias(self): + """The #4531 manager name remains a compatible public alias.""" + assert CheckpointWithoutOutputManager is CheckpointManager + assert isinstance(CheckpointWithoutOutputManager(), CheckpointManager) + + def test_auto_register(self): + """CheckpointWithoutOutput auto-registers to manager when ckpt_manager is provided.""" + manager = CheckpointWithoutOutputManager() + + def func(x): + return x * 2 + 1 + + input_t = torch.randn(4, 4, device='cuda', requires_grad=True) + + ckpt = CheckpointWithoutOutput(ckpt_manager=manager) + y = ckpt.checkpoint(func, input_t) + + assert len(manager.checkpoints) == 1 + assert manager.checkpoints[0] is ckpt + + ckpt2 = CheckpointWithoutOutput(ckpt_manager=manager) + y2 = ckpt2.checkpoint(torch.nn.functional.gelu, y) + + assert len(manager.checkpoints) == 2 + assert manager.checkpoints[1] is ckpt2 + + loss = y2.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss) + loss.backward() + + assert input_t.grad is not None + + def test_discard_is_noop_with_manager(self): + """discard_output_and_register_recompute is a NO-OP when ckpt_manager is set.""" + manager = CheckpointWithoutOutputManager() + + def func1(x): + return x * 2 + + def func2(x): + return torch.nn.functional.gelu(x) + + input_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + y1_ref = func1(input_ref) + y2_ref = func2(y1_ref) + loss_ref = y2_ref.sum() + loss_ref.backward() + grad_ref = input_ref.grad.clone() + + input_ckpt = input_ref.detach().clone().requires_grad_(True) + + ckpt1 = CheckpointWithoutOutput(ckpt_manager=manager) + y1 = ckpt1.checkpoint(func1, input_ckpt) + ckpt1.discard_output_and_register_recompute(y1) + + ckpt2 = CheckpointWithoutOutput(ckpt_manager=manager) + y2 = ckpt2.checkpoint(func2, y1) + ckpt2.discard_output_and_register_recompute(y2) + + assert y1.untyped_storage().size() > 0, "y1 should NOT be discarded yet" + assert y2.untyped_storage().size() > 0, "y2 should NOT be discarded yet" + + loss_ckpt = y2.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + assert y1.untyped_storage().size() == 0, "y1 should be discarded after manager call" + assert y2.untyped_storage().size() == 0, "y2 should be discarded after manager call" + + loss_ckpt.backward() + grad_ckpt = input_ckpt.grad.clone() + + assert torch.allclose(grad_ckpt, grad_ref, atol=1e-6) + + def test_backward_compat_without_manager(self): + """CheckpointWithoutOutput without ckpt_manager should work exactly as before.""" + + def func(x): + return torch.nn.functional.gelu(x) + + input_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + y_ref = func(input_ref) + z_ref = y_ref * 2 + loss_ref = z_ref.sum() + loss_ref.backward() + grad_ref = input_ref.grad.clone() + + input_ckpt = input_ref.detach().clone().requires_grad_(True) + + ckpt = CheckpointWithoutOutput() + y = ckpt.checkpoint(func, input_ckpt) + z = y * 2 + ckpt.discard_output_and_register_recompute(z) + + assert y.untyped_storage().size() == 0 + + loss_ckpt = z.sum() + loss_ckpt.backward() + grad_ckpt = input_ckpt.grad.clone() + + assert torch.allclose(grad_ckpt, grad_ref, atol=1e-6) + + def test_error_handling(self): + """CheckpointWithoutOutputManager rejects invalid add_checkpoint calls.""" + manager = CheckpointWithoutOutputManager() + + with pytest.raises(TypeError): + manager.add_checkpoint("not a checkpoint") + + ckpt = CheckpointWithoutOutput() + with pytest.raises(ValueError): + manager.add_checkpoint(ckpt) + + +class TestCheckpointManagerSequentialChain: + """Test CheckpointWithoutOutputManager with sequential checkpoint chains.""" + + def setup_method(self, method): + Utils.initialize_model_parallel() + initialize_rng_tracker(force_reset=True) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def test_basic_sequential_chain(self): + """Three sequential checkpoints: gradients match non-checkpointed version.""" + + def func1(x): + return x * 2 + 1 + + def func2(x): + return torch.nn.functional.gelu(x) + + def func3(x): + return x * x + x + + input_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + input_ckpt = input_ref.detach().clone().requires_grad_(True) + + y1_ref = func1(input_ref) + y2_ref = func2(y1_ref) + y3_ref = func3(y2_ref) + loss_ref = y3_ref.sum() + loss_ref.backward() + grad_ref = input_ref.grad.clone() + + manager = CheckpointWithoutOutputManager() + + y1 = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func1, input_ckpt) + y2 = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func2, y1) + y3 = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func3, y2) + + loss_ckpt = y3.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + assert y1.untyped_storage().size() == 0, "y1 storage should be released" + assert y2.untyped_storage().size() == 0, "y2 storage should be released" + assert y3.untyped_storage().size() == 0, "y3 storage should be released" + + loss_ckpt.backward() + grad_ckpt = input_ckpt.grad.clone() + + assert torch.allclose( + grad_ckpt, grad_ref, atol=1e-6 + ), f"Gradients mismatch!\nWith manager: {grad_ckpt}\nReference: {grad_ref}" + + def test_sequential_chain_with_dropout(self): + """RNG state is restored during recompute so dropout gradients match.""" + + def func_with_dropout(x): + return torch.nn.functional.dropout(x, p=0.3, training=True) + + def func2(x): + return torch.nn.functional.gelu(x) + + input_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + input_ckpt = input_ref.detach().clone().requires_grad_(True) + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + y1_ref = func_with_dropout(input_ref) + y2_ref = func2(y1_ref) + loss_ref = y2_ref.sum() + loss_ref.backward() + grad_ref = input_ref.grad.clone() + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + manager = CheckpointWithoutOutputManager() + + y1 = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func_with_dropout, input_ckpt) + y2 = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func2, y1) + + loss_ckpt = y2.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + loss_ckpt.backward() + grad_ckpt = input_ckpt.grad.clone() + + assert torch.allclose( + grad_ckpt, grad_ref, atol=1e-6 + ), f"Gradients with dropout mismatch!\nWith manager: {grad_ckpt}\nReference: {grad_ref}" + + def test_multiple_outputs(self): + """CheckpointWithoutOutputManager handles functions that return multiple outputs.""" + + def func_multi_output(x): + return x * 2, x + 1 + + def func_combine(a, b): + return a + b + + input_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + input_ckpt = input_ref.detach().clone().requires_grad_(True) + + y1a_ref, y1b_ref = func_multi_output(input_ref) + y2_ref = func_combine(y1a_ref, y1b_ref) + loss_ref = y2_ref.sum() + loss_ref.backward() + grad_ref = input_ref.grad.clone() + + manager = CheckpointWithoutOutputManager() + + y1a, y1b = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint( + func_multi_output, input_ckpt + ) + y2 = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func_combine, y1a, y1b) + + loss_ckpt = y2.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + loss_ckpt.backward() + grad_ckpt = input_ckpt.grad.clone() + + assert torch.allclose(grad_ckpt, grad_ref, atol=1e-6), ( + f"Gradients mismatch with multiple outputs!\n" + f"With manager: {grad_ckpt}\nReference: {grad_ref}" + ) + + +class TestCheckpointManagerPartialCheckpoint: + """Test CheckpointWithoutOutputManager with partial checkpointing (some ops not checkpointed).""" + + def setup_method(self, method): + Utils.initialize_model_parallel() + initialize_rng_tracker(force_reset=True) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def test_partial_checkpoint(self): + """ + Only f and h are checkpointed; g is a regular operation. + + Computation chain: + a --[f]--> b --[g]--> c --[h]--> d --[sum]--> loss + """ + + def func_f(x): + return torch.nn.functional.gelu(x * 2 + 1) + + def func_g(x): + return x * 3 - 2 + + def func_h(x): + return torch.sigmoid(x) + x + + input_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + + b_ref = func_f(input_ref) + c_ref = func_g(b_ref) + d_ref = func_h(c_ref) + loss_ref = d_ref.sum() + loss_ref.backward() + grad_ref = input_ref.grad.clone() + + input_ckpt = input_ref.detach().clone().requires_grad_(True) + + manager = CheckpointWithoutOutputManager() + + b = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func_f, input_ckpt) + c = func_g(b) + d = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func_h, c) + + loss_ckpt = d.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + assert b.untyped_storage().size() == 0, "b storage should be released" + assert d.untyped_storage().size() == 0, "d storage should be released" + assert c.untyped_storage().size() > 0, "c storage should NOT be released (not checkpointed)" + + loss_ckpt.backward() + grad_ckpt = input_ckpt.grad.clone() + + assert torch.allclose(grad_ckpt, grad_ref, atol=1e-6), ( + f"Gradients mismatch with partial checkpoint!\n" + f"With manager: {grad_ckpt}\nReference: {grad_ref}" + ) + + def test_partial_checkpoint_with_tuple_output(self): + """ + Mimics HyperConnection's computation pattern with tuple outputs. + + - compute_mappings: checkpointed, returns tuple (h_pre, h_post, h_res) + - aggregate: NOT checkpointed + - apply_h_res: checkpointed + - apply_h_post: checkpointed + """ + + def compute_mappings(x): + h_pre = torch.sigmoid(x.mean(dim=-1, keepdim=True).expand_as(x)) + h_post = torch.tanh(x.sum(dim=-1, keepdim=True).expand_as(x)) + h_res = torch.relu(x) + return h_pre, h_post, h_res + + def aggregate(x, h_pre): + return x * h_pre + + def apply_h_res(h_res, residual): + return h_res + residual * 0.5 + + def apply_h_post(y, h_post): + return y * h_post + y + + x_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + residual_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + + h_pre_ref, h_post_ref, h_res_ref = compute_mappings(x_ref) + agg_ref = aggregate(x_ref, h_pre_ref) + y_ref = torch.nn.functional.gelu(agg_ref) + mixed_ref = apply_h_res(h_res_ref, residual_ref) + output_ref = apply_h_post(y_ref, h_post_ref) + final_ref = output_ref + mixed_ref + loss_ref = final_ref.sum() + loss_ref.backward() + grad_x_ref = x_ref.grad.clone() + grad_residual_ref = residual_ref.grad.clone() + + x_ckpt = x_ref.detach().clone().requires_grad_(True) + residual_ckpt = residual_ref.detach().clone().requires_grad_(True) + + manager = CheckpointWithoutOutputManager() + + h_pre, h_post, h_res = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint( + compute_mappings, x_ckpt + ) + agg = aggregate(x_ckpt, h_pre) + y = torch.nn.functional.gelu(agg) + mixed = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint( + apply_h_res, h_res, residual_ckpt + ) + output = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(apply_h_post, y, h_post) + + final = output + mixed + loss_ckpt = final.sum() + + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + assert h_pre.untyped_storage().size() == 0, "h_pre storage should be released" + assert h_post.untyped_storage().size() == 0, "h_post storage should be released" + assert h_res.untyped_storage().size() == 0, "h_res storage should be released" + assert mixed.untyped_storage().size() == 0, "mixed storage should be released" + assert output.untyped_storage().size() == 0, "output storage should be released" + + assert agg.untyped_storage().size() > 0, "agg storage should NOT be released" + assert y.untyped_storage().size() > 0, "y storage should NOT be released" + + loss_ckpt.backward() + grad_x_ckpt = x_ckpt.grad.clone() + grad_residual_ckpt = residual_ckpt.grad.clone() + + assert torch.allclose( + grad_x_ckpt, grad_x_ref, atol=1e-6 + ), f"Gradients for x mismatch!\nWith manager: {grad_x_ckpt}\nReference: {grad_x_ref}" + assert torch.allclose(grad_residual_ckpt, grad_residual_ref, atol=1e-6), ( + f"Gradients for residual mismatch!\n" + f"With manager: {grad_residual_ckpt}\nReference: {grad_residual_ref}" + ) + + +# ============================================================================ +# Block-level mHC recompute coverage +# ============================================================================ +# +# These tests instantiate a full ``TransformerBlock`` with mHC enabled to +# exercise: +# * ``_build_mhc_recompute_layer_plan`` (per-layer ``CheckpointWithoutOutputManager`` +# allocation, including the ``mhc_recompute_layer_num`` boundary case), +# * ``_finalize_mhc_recompute_layer`` (manager finalization at block end), +# * the ``HyperConnectionModule.input_expand`` / ``output_contract`` calls +# in ``TransformerBlock.forward`` for ``pre_process`` / ``post_process`` +# stages. +# +# Single-process (no PP) so they can run on a single-GPU CI lane. + + +class TestTransformerBlockMHCRecompute: + """End-to-end ``TransformerBlock`` forward with mHC selective recompute.""" + + def setup_method(self, method): + from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + @staticmethod + def _make_mhc_block(num_layers, num_streams=4, mhc_recompute_layer_num=None): + from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_layer_with_transformer_engine_spec, + ) + from megatron.core.transformer.hyper_connection import HyperConnectionModule + from megatron.core.transformer.transformer_block import TransformerBlock + from megatron.core.transformer.transformer_config import TransformerConfig + from megatron.core.transformer.transformer_layer import HyperConnectionTransformerLayer + + config = TransformerConfig( + num_layers=num_layers, + hidden_size=64, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=num_streams, + mhc_sinkhorn_iterations=5, + mhc_init_gating_factor=0.01, + mhc_recompute_layer_num=mhc_recompute_layer_num, + recompute_granularity='selective', + recompute_modules=['mhc'], + hidden_dropout=0.0, + attention_dropout=0.0, + ) + spec = get_gpt_layer_with_transformer_engine_spec() + spec.module = HyperConnectionTransformerLayer + spec.submodules.self_attention_hyper_connection = HyperConnectionModule + spec.submodules.mlp_hyper_connection = HyperConnectionModule + return TransformerBlock(config, spec, pre_process=True, post_process=True).cuda(), config + + def _check_recompute_plan(self, block, expected_block_ends): + """Drive ``_build_mhc_recompute_layer_plan`` directly and check the boundary list.""" + block.train() + managers, ends = block._build_mhc_recompute_layer_plan(use_mhc_recompute=True) + assert len(managers) == len(block.layers) + assert ends == expected_block_ends, f"got {ends}, expected {expected_block_ends}" + # Layers in the same recompute block share a manager; new block → new manager. + last_was_end = True + last_mgr = None + for mgr, end in zip(managers, ends): + assert mgr is not None + if last_was_end: + assert mgr is not last_mgr, "new recompute block should get a new manager" + else: + assert mgr is last_mgr, "layers within a recompute block share a manager" + last_was_end = end + last_mgr = mgr + + def test_recompute_plan_no_layer_num(self): + """Without ``mhc_recompute_layer_num`` only the final layer ends a recompute block.""" + block, _ = self._make_mhc_block(num_layers=4) + self._check_recompute_plan(block, expected_block_ends=[False, False, False, True]) + + def test_recompute_plan_with_layer_num(self): + """With ``mhc_recompute_layer_num=2`` every other layer ends a recompute block.""" + block, _ = self._make_mhc_block(num_layers=4, mhc_recompute_layer_num=2) + self._check_recompute_plan(block, expected_block_ends=[False, True, False, True]) + + def test_recompute_plan_disabled(self): + """``use_mhc_recompute=False`` returns an all-None / all-False plan.""" + block, _ = self._make_mhc_block(num_layers=3) + managers, ends = block._build_mhc_recompute_layer_plan(use_mhc_recompute=False) + assert managers == [None, None, None] + assert ends == [False, False, False] + + def test_block_forward_input_expand_output_contract(self): + """Forward exercises ``input_expand`` (pre) and ``output_contract`` (post).""" + block, config = self._make_mhc_block(num_layers=2, mhc_recompute_layer_num=2) + block.train() + + seq_len = 8 + batch_size = 2 + # Input is [s, b, hidden_size]; the block must expand to [s, b, n*hidden_size] + # internally, then contract back to [s, b, hidden_size] before final layernorm. + hidden_states = torch.randn( + seq_len, batch_size, config.hidden_size, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=torch.bool, device='cuda') + + out = block(hidden_states=hidden_states, attention_mask=attention_mask) + assert out.shape == hidden_states.shape, ( + f"output_contract should restore original shape, got {tuple(out.shape)} " + f"vs expected {tuple(hidden_states.shape)}" + ) + # Backward should flow through the recompute path without error. + out.sum().backward() + assert hidden_states.grad is not None + assert torch.isfinite(hidden_states.grad).all() diff --git a/tests/unit_tests/transformer/test_module.py b/tests/unit_tests/transformer/test_module.py index 92f15b2f46d..5faf6c81ef1 100644 --- a/tests/unit_tests/transformer/test_module.py +++ b/tests/unit_tests/transformer/test_module.py @@ -4,7 +4,7 @@ import torch from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed -from megatron.core.transformer.module import Float16Module, MegatronModule +from megatron.core.transformer.module import Float16Module, MegatronModule, mark_keep_in_fp32 from megatron.core.transformer.transformer_config import TransformerConfig from tests.unit_tests.test_utilities import Utils @@ -163,3 +163,18 @@ def test_bf16_module(self): x = torch.ones((2, 2)).cuda() # inputs are converted to bf16 then outputs are converted to fp32 assert bf16_module(x).dtype == torch.float32 + + @pytest.mark.parametrize( + ('precision', 'dtype'), [('fp16', torch.float16), ('bf16', torch.bfloat16)] + ) + def test_keep_in_fp32_params(self, precision, dtype): + transformer_config = self.transformer_config + megatron_module = self.megatron_module + megatron_module.fp32_param = mark_keep_in_fp32( + torch.nn.Parameter(torch.zeros(4, dtype=torch.float32, device='cuda')) + ) + setattr(transformer_config, precision, True) + float16_module = Float16Module(config=transformer_config, module=megatron_module) + + assert float16_module.module.linear.weight.dtype == dtype + assert float16_module.module.fp32_param.dtype == torch.float32 diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py index c3c3944e007..59a4d5cd887 100644 --- a/tests/unit_tests/transformer/test_multi_token_prediction.py +++ b/tests/unit_tests/transformer/test_multi_token_prediction.py @@ -22,6 +22,7 @@ from megatron.core.parallel_state import get_context_parallel_group from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.hyper_connection import learned_output_contract from megatron.core.transformer.multi_token_prediction import ( MTPLossLoggingHelper, MultiTokenPredictionBlock, @@ -168,7 +169,10 @@ def test_constructor_ues_te(self, tp, cp): assert num_weights == 15216 * config.mtp_num_layers def test_get_embeddings_rolls_padding_mask(self): - """Test that _get_embeddings rolls padding_mask alongside input ids.""" + """Test that _get_embeddings rolls padding_mask alongside input ids. + + The padding mask uses True for padded positions, including new MTP boundaries. + """ torch.manual_seed(_SEED) config, mtp_block_spec = self._create_config_and_mtp_block_spec(tp=1, cp=1) mtp = MultiTokenPredictionBlock(config=config, spec=mtp_block_spec) @@ -179,7 +183,7 @@ def test_get_embeddings_rolls_padding_mask(self): input_ids = torch.tensor([[1, 2, 3, 4, 0, 0], [5, 6, 7, 0, 0, 0]], dtype=torch.int64) position_ids = torch.arange(seq_len, dtype=torch.int64).repeat(batch_size, 1) padding_mask = torch.tensor( - [[True, True, True, True, False, False], [True, True, True, False, False, False]] + [[False, False, False, False, True, True], [False, False, False, True, True, True]] ) hidden_states = torch.randn(seq_len, batch_size, config.hidden_size) @@ -199,14 +203,14 @@ def fake_embedding(input_ids, position_ids): expected_input_ids, _ = roll_tensor(input_ids, shifts=-1, dims=-1) expected_position_ids, _ = roll_tensor(position_ids, shifts=-1, dims=-1) - expected_padding_mask, _ = roll_tensor(padding_mask, shifts=-1, dims=-1) + expected_padding_mask, _ = roll_tensor(padding_mask, shifts=-1, dims=-1, fill_value=True) assert torch.equal(rolled_input_ids, expected_input_ids) assert torch.equal(rolled_position_ids, expected_position_ids) assert torch.equal(rolled_padding_mask, expected_padding_mask) def test_forward_propagates_rolled_padding_mask(self, monkeypatch): - """Test forward passes rolled padding_mask to transformer path.""" + """Test forward passes the boundary-padded rolled mask to the transformer path.""" torch.manual_seed(_SEED) config, mtp_block_spec = self._create_config_and_mtp_block_spec(tp=1, cp=1) mtp = MultiTokenPredictionBlock(config=config, spec=mtp_block_spec) @@ -216,7 +220,7 @@ def test_forward_propagates_rolled_padding_mask(self, monkeypatch): batch_size = 2 input_ids = torch.tensor([[1, 2, 3, 0], [4, 5, 0, 0]], dtype=torch.int64) position_ids = torch.arange(seq_len, dtype=torch.int64).repeat(batch_size, 1) - padding_mask = torch.tensor([[True, True, True, False], [True, True, False, False]]) + padding_mask = torch.tensor([[False, False, False, True], [False, False, True, True]]) hidden_states = torch.randn(seq_len, batch_size, config.hidden_size) attention_mask = torch.ones((batch_size, 1, seq_len, seq_len), dtype=torch.bool) seen = {} @@ -258,7 +262,7 @@ def fake_proj_and_transformer_layer( embedding=fake_embedding, ) - expected_padding_mask, _ = roll_tensor(padding_mask, shifts=-1, dims=-1) + expected_padding_mask, _ = roll_tensor(padding_mask, shifts=-1, dims=-1, fill_value=True) assert torch.equal(seen["padding_mask"], expected_padding_mask) assert torch.equal(returned_padding_mask, expected_padding_mask) @@ -1080,6 +1084,179 @@ def test_roll_tensor_with_packed_sequences(self, cp): Utils.destroy_model_parallel() + def test_roll_tensor_with_packed_sequences_contiguous_cp(self): + """Contiguous THD CP rolls across rank boundaries without crossing sequence boundaries.""" + cp = 2 + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=cp) + cp_group = get_context_parallel_group() + cp_rank = torch.distributed.get_rank(group=cp_group) + + # Full padded layout: + # seq1: [1,2,3,4,5,6,7,0] + # seq2: [11,12,13,14,15,16,17,18,19,20,21,0] + # Contiguous CP rank 0 owns global rows [0, 10), rank 1 owns [10, 20). + if cp_rank == 0: + tensor = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 0, 11, 12]], dtype=torch.float32).cuda() + expected = torch.tensor([[2, 3, 4, 5, 6, 7, 0, 0, 12, 13]], dtype=torch.float32).cuda() + padding_mask = torch.tensor( + [[False, False, False, False, False, False, False, True, False, False]] + ).cuda() + expected_padding_mask = torch.tensor( + [[False, False, False, False, False, False, True, True, False, False]] + ).cuda() + else: + tensor = torch.tensor( + [[13, 14, 15, 16, 17, 18, 19, 20, 21, 0]], dtype=torch.float32 + ).cuda() + expected = torch.tensor( + [[14, 15, 16, 17, 18, 19, 20, 21, 0, 0]], dtype=torch.float32 + ).cuda() + padding_mask = torch.tensor( + [[False, False, False, False, False, False, False, False, False, True]] + ).cuda() + expected_padding_mask = torch.tensor( + [[False, False, False, False, False, False, False, False, True, True]] + ).cuda() + + cu_seqlens = torch.tensor([0, 7, 18], dtype=torch.int32).cuda() + cu_seqlens_padded = torch.tensor([0, 8, 20], dtype=torch.int32).cuda() + packed_seq_params = PackedSeqParams( + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=11, + max_seqlen_kv=11, + qkv_format='thd', + ) + packed_seq_params.cp_partition_mode = 'contiguous' + + rolled, sum_val = roll_tensor( + tensor, + shifts=-1, + dims=-1, + cp_group=cp_group, + packed_seq_params=packed_seq_params, + ) + rolled_padding_mask, _ = roll_tensor( + padding_mask, + shifts=-1, + dims=-1, + cp_group=cp_group, + packed_seq_params=packed_seq_params, + fill_value=True, + ) + + assert torch.equal(rolled, expected), ( + f"CP Rank {cp_rank}: Expected\n{expected}\nbut got\n{rolled}\nDiff:\n" + f"{rolled - expected}" + ) + assert torch.equal(rolled_padding_mask, expected_padding_mask), ( + f"CP Rank {cp_rank}: Expected padding mask\n{expected_padding_mask}\nbut got\n" + f"{rolled_padding_mask}" + ) + assert sum_val.numel() == 1, "Sum should be a scalar" + + Utils.destroy_model_parallel() + + @pytest.mark.parametrize("cp", [1, 2]) + def test_roll_tensor_with_packed_sequences_odd_seqlen(self, cp): + """Test roll_tensor with ODD packed seqlens. + + For CP=1: per-sequence rolling on contiguous packed tensor — odd seqlens are fine + with cu_seqlens_q alone (no padding required). + For CP=2: each per-sequence padded length must be a multiple of 2*cp_size, so odd + seqlens require padding. The local THD-CP layout is determined by + cu_seqlens_q_padded; the roll function must use the padded boundaries to + index local chunks correctly. Without the padded boundaries, real tokens + leak across sequence boundaries. + """ + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=cp) + cp_group = get_context_parallel_group() if cp > 1 else None + cp_rank = torch.distributed.get_rank(group=cp_group) if cp_group is not None else 0 + + if cp == 1: + # Two odd-length sequences: [3, 5]. Total = 8. + tensor = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.float32).cuda() + cu_seqlens = torch.tensor([0, 3, 8], dtype=torch.int32).cuda() + + packed_seq_params = PackedSeqParams( + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=5, + max_seqlen_kv=5, + qkv_format='thd', + ) + + rolled, sum_val = roll_tensor( + tensor, shifts=-1, dims=0, cp_group=cp_group, packed_seq_params=packed_seq_params + ) + + # seq1 [1,2,3] -> [2,3,0]; seq2 [4,5,6,7,8] -> [5,6,7,8,0] + expected = torch.tensor([2, 3, 0, 5, 6, 7, 8, 0], dtype=torch.float32).cuda() + assert torch.equal(rolled, expected), f"Expected {expected}, got {rolled}" + else: + # Two ODD sequences padded up to multiples of 2*cp_size = 4: + # seq1: real=[1..7] (len 7), padded with 0 -> [1,2,3,4,5,6,7,0] (len 8) + # seq2: real=[11..21] (len 11), padded with 0 -> + # [11,12,13,14,15,16,17,18,19,20,21,0] (len 12) + # Zigzag (4 chunks per padded seq, rank r owns chunks (r, 3-r)): + # seq1 chunks: [1,2], [3,4], [5,6], [7,0] + # rank 0 -> [1,2, 7,0]; rank 1 -> [3,4, 5,6] + # seq2 chunks: [11,12,13], [14,15,16], [17,18,19], [20,21,0] + # rank 0 -> [11,12,13, 20,21,0]; rank 1 -> [14,15,16, 17,18,19] + # Expected after roll(-1) within unpadded region (last real -> 0; pad stays 0): + # seq1 rolled real: [2,3,4,5,6,7,0]; padded last -> 0 + # seq2 rolled real: [12,13,14,15,16,17,18,19,20,21,0]; padded last -> 0 + # Re-zigzag the rolled+padded seqs: + # seq1: [2,3], [4,5], [6,7], [0,0] + # rank 0 -> [2,3, 0,0]; rank 1 -> [4,5, 6,7] + # seq2: [12,13,14], [15,16,17], [18,19,20], [21,0,0] + # rank 0 -> [12,13,14, 21,0,0]; rank 1 -> [15,16,17, 18,19,20] + if cp_rank == 0: + tensor = torch.tensor( + [1, 2, 7, 0, 11, 12, 13, 20, 21, 0], dtype=torch.float32 + ).cuda() + expected = torch.tensor( + [2, 3, 0, 0, 12, 13, 14, 21, 0, 0], dtype=torch.float32 + ).cuda() + else: + tensor = torch.tensor( + [3, 4, 5, 6, 14, 15, 16, 17, 18, 19], dtype=torch.float32 + ).cuda() + expected = torch.tensor( + [4, 5, 6, 7, 15, 16, 17, 18, 19, 20], dtype=torch.float32 + ).cuda() + + # Unpadded cu_seqlens_q = [0, 7, 18]; padded = [0, 8, 20]. + cu_seqlens = torch.tensor([0, 7, 18], dtype=torch.int32).cuda() + cu_seqlens_padded = torch.tensor([0, 8, 20], dtype=torch.int32).cuda() + + packed_seq_params = PackedSeqParams( + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=11, + max_seqlen_kv=11, + qkv_format='thd', + ) + + rolled, sum_val = roll_tensor( + tensor, shifts=-1, dims=0, cp_group=cp_group, packed_seq_params=packed_seq_params + ) + + assert ( + rolled.shape == expected.shape + ), f"Shape mismatch: expected {expected.shape}, got {rolled.shape}" + assert torch.equal( + rolled, expected + ), f"CP Rank {cp_rank}: Expected\n{expected}\nbut got\n{rolled}\nDiff:\n{rolled - expected}" + + assert sum_val.numel() == 1, "Sum should be a scalar" + + Utils.destroy_model_parallel() + class TestMTPLossLoggingHelper: def setup_method(self, method): @@ -1528,3 +1705,52 @@ def test_attention_mask_validation_mamba(self): pytest.fail(f"Attention mask validation failed for Mamba hybrid model: {e}") else: raise + + +class TestLearnedOutputContract: + """Tests for the learned n-stream to one-stream mHC contraction.""" + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(_SEED) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) + def test_shape_and_dtype(self, dtype): + hidden_size, n_streams = 32, 4 + hidden_states = torch.randn(8, 2, n_streams * hidden_size, device="cuda", dtype=dtype) + head_fn = torch.randn(n_streams, n_streams * hidden_size, device="cuda") + base = torch.zeros(n_streams, device="cuda") + scale = torch.ones(1, device="cuda") + + output = learned_output_contract(hidden_states, head_fn, base, scale, n_streams, eps=1e-6) + + assert output.shape == (8, 2, hidden_size) + assert output.dtype == dtype + + def test_gradient_and_reference(self): + hidden_size, n_streams, eps = 8, 2, 1e-6 + hidden_states = torch.randn( + 2, 1, n_streams * hidden_size, device="cuda", dtype=torch.float32, requires_grad=True + ) + head_fn = torch.randn(n_streams, n_streams * hidden_size, device="cuda", requires_grad=True) + base = torch.zeros(n_streams, device="cuda", requires_grad=True) + scale = torch.ones(1, device="cuda", requires_grad=True) + + output = learned_output_contract(hidden_states, head_fn, base, scale, n_streams, eps) + rsqrt = torch.rsqrt(hidden_states.square().mean(-1, keepdim=True) + eps) + mixes = torch.nn.functional.linear(hidden_states, head_fn) * rsqrt + weights = torch.sigmoid(mixes * scale + base) + eps + expected = torch.sum( + weights.unsqueeze(-1) + * hidden_states.view(*hidden_states.shape[:-1], n_streams, hidden_size), + dim=-2, + ) + torch.testing.assert_close(output, expected) + + output.sum().backward() + for tensor in (hidden_states, head_fn, base, scale): + assert tensor.grad is not None + assert torch.count_nonzero(tensor.grad) > 0 diff --git a/tests/unit_tests/transformer/test_packed_seq_params_cuda_graph.py b/tests/unit_tests/transformer/test_packed_seq_params_cuda_graph.py new file mode 100644 index 00000000000..f838c825c37 --- /dev/null +++ b/tests/unit_tests/transformer/test_packed_seq_params_cuda_graph.py @@ -0,0 +1,851 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import pytest +import torch + +from megatron.core.packed_seq_params import ( + CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX, + PACKED_SEQ_PARAMS_CUDA_GRAPH_STATIC_FIELDS, + PACKED_SEQ_PARAMS_CUDA_GRAPH_TENSOR_FIELDS, + PackedSeqParams, + build_packed_seq_params_from_cuda_graph_kwargs, + has_packed_seq_params_cuda_graph_kwargs, + split_packed_seq_params_for_cuda_graph, +) +from megatron.core.transformer.cuda_graphs import ( + _add_packed_seq_params_to_te_cuda_graph_sample_kwargs, +) +from megatron.core.transformer.transformer_layer import TransformerLayer + + +class _TransformerLayerCudaGraphStub: + _set_te_cuda_graph_packed_seq_params_static_metadata = ( + TransformerLayer._set_te_cuda_graph_packed_seq_params_static_metadata + ) + _get_te_cuda_graph_packed_seq_params_static_metadata = ( + TransformerLayer._get_te_cuda_graph_packed_seq_params_static_metadata + ) + _validate_te_cuda_graph_packed_seq_params_static_metadata = ( + TransformerLayer._validate_te_cuda_graph_packed_seq_params_static_metadata + ) + _get_te_cuda_graph_packed_seq_params_tensor_kwarg_names = ( + TransformerLayer._get_te_cuda_graph_packed_seq_params_tensor_kwarg_names + ) + _validate_te_cuda_graph_packed_seq_params_tensor_kwargs = ( + TransformerLayer._validate_te_cuda_graph_packed_seq_params_tensor_kwargs + ) + _rebuild_te_cuda_graph_packed_seq_params = ( + TransformerLayer._rebuild_te_cuda_graph_packed_seq_params + ) + _flatten_te_cuda_graph_packed_seq_params = ( + TransformerLayer._flatten_te_cuda_graph_packed_seq_params + ) + + +def _make_packed_seq_params(): + cu_seqlens = torch.IntTensor([0, 4, 9, 16]) + cu_seqlens_padded = torch.IntTensor([0, 8, 12, 16]) + return PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=8, + max_seqlen_kv=8, + local_cp_size=1, + pad_between_seqs=False, + cp_partition_mode="contiguous", + ) + + +def test_split_packed_seq_params_for_cuda_graph_separates_tensors_from_metadata(): + packed_seq_params = _make_packed_seq_params() + + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params) + + assert set(static_metadata) == set(PACKED_SEQ_PARAMS_CUDA_GRAPH_STATIC_FIELDS) + assert static_metadata == { + "qkv_format": "thd", + "max_seqlen_q": 8, + "max_seqlen_kv": 8, + "local_cp_size": 1, + "cp_group": None, + "pad_between_seqs": False, + "cp_partition_mode": "contiguous", + } + assert all(not isinstance(value, torch.Tensor) for value in static_metadata.values()) + + expected_tensor_fields = { + "cu_seqlens_q", + "cu_seqlens_kv", + "cu_seqlens_q_padded", + "cu_seqlens_kv_padded", + } + assert set(tensor_kwargs) == { + f"{CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX}{field}" for field in expected_tensor_fields + } + assert set(PACKED_SEQ_PARAMS_CUDA_GRAPH_TENSOR_FIELDS) >= expected_tensor_fields + for value in tensor_kwargs.values(): + assert isinstance(value, torch.Tensor) + + +def test_has_packed_seq_params_cuda_graph_kwargs_detects_flattened_fields(): + tensor_kwargs, _ = split_packed_seq_params_for_cuda_graph(_make_packed_seq_params()) + + assert has_packed_seq_params_cuda_graph_kwargs(tensor_kwargs) + assert not has_packed_seq_params_cuda_graph_kwargs({"hidden_states": torch.ones(2, 1, 4)}) + assert build_packed_seq_params_from_cuda_graph_kwargs({}, None) is None + + +def test_build_packed_seq_params_from_cuda_graph_kwargs_pops_flattened_fields(): + packed_seq_params = _make_packed_seq_params() + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params) + kwargs = {"hidden_states": torch.ones(2, 1, 4), **tensor_kwargs} + + rebuilt = build_packed_seq_params_from_cuda_graph_kwargs(kwargs, static_metadata) + + assert set(kwargs) == {"hidden_states"} + assert rebuilt.qkv_format == "thd" + assert rebuilt.max_seqlen_q == 8 + assert rebuilt.max_seqlen_kv == 8 + assert rebuilt.local_cp_size == 1 + assert rebuilt.cp_group is None + assert rebuilt.pad_between_seqs is False + assert rebuilt.cp_partition_mode == "contiguous" + assert rebuilt.total_tokens is None + assert rebuilt.seq_idx is None + assert torch.equal(rebuilt.cu_seqlens_q, packed_seq_params.cu_seqlens_q) + assert torch.equal(rebuilt.cu_seqlens_kv, packed_seq_params.cu_seqlens_kv) + assert torch.equal(rebuilt.cu_seqlens_q_padded, packed_seq_params.cu_seqlens_q_padded) + assert torch.equal(rebuilt.cu_seqlens_kv_padded, packed_seq_params.cu_seqlens_kv_padded) + + +def test_build_packed_seq_params_from_cuda_graph_kwargs_can_keep_kwargs_intact(): + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph( + _make_packed_seq_params() + ) + kwargs = dict(tensor_kwargs) + + build_packed_seq_params_from_cuda_graph_kwargs( + kwargs, static_metadata, remove_from_kwargs=False + ) + + assert kwargs == tensor_kwargs + + +def test_split_packed_seq_params_for_cuda_graph_rejects_static_tensor_metadata(): + packed_seq_params = _make_packed_seq_params() + packed_seq_params.max_seqlen_q = torch.IntTensor([8]) + + with pytest.raises(TypeError, match="max_seqlen_q"): + split_packed_seq_params_for_cuda_graph(packed_seq_params) + + +def test_split_packed_seq_params_for_cuda_graph_ignores_mamba_only_fields(): + packed_seq_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=torch.IntTensor([0, 2, 5]), + cu_seqlens_kv=torch.IntTensor([0, 2, 5]), + max_seqlen_q=3, + max_seqlen_kv=3, + total_tokens=5, + ) + assert packed_seq_params.seq_idx is not None + + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params) + + assert f"{CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX}seq_idx" not in tensor_kwargs + assert "total_tokens" not in static_metadata + + +def test_transformer_layer_rebuilds_flattened_cuda_graph_packed_seq_params(): + layer = _TransformerLayerCudaGraphStub() + packed_seq_params = _make_packed_seq_params() + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params) + layer._set_te_cuda_graph_packed_seq_params_static_metadata(static_metadata, tensor_kwargs) + kwargs = {"hidden_states": torch.ones(2, 1, 4), **tensor_kwargs} + + layer._rebuild_te_cuda_graph_packed_seq_params(kwargs) + + assert set(kwargs) == {"hidden_states", "packed_seq_params"} + rebuilt = kwargs["packed_seq_params"] + assert rebuilt.qkv_format == "thd" + assert rebuilt.max_seqlen_q == 8 + assert rebuilt.max_seqlen_kv == 8 + assert rebuilt.cp_partition_mode == "contiguous" + assert torch.equal(rebuilt.cu_seqlens_q, packed_seq_params.cu_seqlens_q) + assert torch.equal(rebuilt.cu_seqlens_kv, packed_seq_params.cu_seqlens_kv) + + +def test_transformer_layer_flattens_replay_time_packed_seq_params(): + layer = _TransformerLayerCudaGraphStub() + packed_seq_params = _make_packed_seq_params() + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params) + layer._set_te_cuda_graph_packed_seq_params_static_metadata(static_metadata, tensor_kwargs) + attention_mask = torch.zeros(1, 1, 16, 16, dtype=torch.bool) + kwargs = {"attention_mask": attention_mask, "packed_seq_params": packed_seq_params} + + layer._flatten_te_cuda_graph_packed_seq_params(kwargs) + + assert kwargs["attention_mask"] is attention_mask + assert "packed_seq_params" not in kwargs + assert set(tensor_kwargs).issubset(kwargs) + for key, value in tensor_kwargs.items(): + assert kwargs[key] is value + + +def test_transformer_layer_rejects_replay_without_captured_packed_seq_params(): + layer = _TransformerLayerCudaGraphStub() + _, static_metadata = split_packed_seq_params_for_cuda_graph(_make_packed_seq_params()) + layer._set_te_cuda_graph_packed_seq_params_static_metadata(static_metadata) + + with pytest.raises(AssertionError, match="captured with packed_seq_params"): + layer._flatten_te_cuda_graph_packed_seq_params({"hidden_states": torch.ones(2, 1, 4)}) + + +def test_transformer_layer_rejects_changed_packed_seq_params_static_metadata(): + layer = _TransformerLayerCudaGraphStub() + packed_seq_params = _make_packed_seq_params() + _, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params) + layer._set_te_cuda_graph_packed_seq_params_static_metadata(static_metadata) + packed_seq_params.max_seqlen_q = 4 + + with pytest.raises(AssertionError, match="max_seqlen_q"): + layer._flatten_te_cuda_graph_packed_seq_params({"packed_seq_params": packed_seq_params}) + + +def test_transformer_layer_rejects_changed_pad_between_seqs_metadata(): + layer = _TransformerLayerCudaGraphStub() + packed_seq_params = _make_packed_seq_params() + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph( + packed_seq_params + ) + layer._set_te_cuda_graph_packed_seq_params_static_metadata( + static_metadata, tensor_kwargs + ) + packed_seq_params.pad_between_seqs = True + + with pytest.raises(AssertionError, match="pad_between_seqs"): + layer._flatten_te_cuda_graph_packed_seq_params( + {"packed_seq_params": packed_seq_params} + ) + + +def test_transformer_layer_rejects_changed_cp_partition_mode_metadata(): + layer = _TransformerLayerCudaGraphStub() + packed_seq_params = _make_packed_seq_params() + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph( + packed_seq_params + ) + layer._set_te_cuda_graph_packed_seq_params_static_metadata( + static_metadata, tensor_kwargs + ) + packed_seq_params.cp_partition_mode = "zigzag" + + with pytest.raises(AssertionError, match="cp_partition_mode"): + layer._flatten_te_cuda_graph_packed_seq_params( + {"packed_seq_params": packed_seq_params} + ) + + +def test_hybrid_wrapper_delegates_prefixed_packed_sequence_contract(): + from megatron.core.models.hybrid.hybrid_block import HyperConnectionHybridLayer + + inner = _TransformerLayerCudaGraphStub() + wrapper = object.__new__(HyperConnectionHybridLayer) + object.__setattr__(wrapper, "inner_layer", inner) + packed_seq_params = _make_packed_seq_params() + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph( + packed_seq_params + ) + + wrapper._set_te_cuda_graph_packed_seq_params_static_metadata( + static_metadata, tensor_kwargs + ) + replay_kwargs = {"packed_seq_params": packed_seq_params} + wrapper._flatten_te_cuda_graph_packed_seq_params(replay_kwargs) + + assert "packed_seq_params" not in replay_kwargs + assert set(replay_kwargs) == set(tensor_kwargs) + wrapper._rebuild_te_cuda_graph_packed_seq_params(replay_kwargs) + assert replay_kwargs["packed_seq_params"].pad_between_seqs is False + assert torch.equal( + replay_kwargs["packed_seq_params"].cu_seqlens_q, + packed_seq_params.cu_seqlens_q, + ) + + +def test_thd_graph_discovery_excludes_wrapped_mamba_without_metadata_contract(): + from types import SimpleNamespace + + from megatron.core.models.hybrid.hybrid_block import HyperConnectionHybridLayer + from megatron.core.ssm.mamba_layer import MambaLayer + from megatron.core.transformer.cuda_graphs import _layer_is_graphable + from megatron.core.transformer.enums import CudaGraphModule + + wrapper = object.__new__(HyperConnectionHybridLayer) + object.__setattr__(wrapper, "inner_layer", object.__new__(MambaLayer)) + config = SimpleNamespace( + cuda_graph_modules=[CudaGraphModule.mamba], + sequence_packing_scheduler="dp_balanced", + ) + + assert not _layer_is_graphable(wrapper, config) + + +def test_hybrid_wrapper_leaves_local_cudagraph_manager_on_inner_layer(): + from megatron.core.models.hybrid.hybrid_block import HyperConnectionHybridLayer + from megatron.core.transformer.module import GraphableMegatronModule + from megatron.core.transformer.transformer_config import TransformerConfig + + class _LocalGraphInner(GraphableMegatronModule): + def __init__(self, config): + super().__init__(config) + self.layer_number = 1 + self.offload_module_in_cuda_graph = True + + def create_mcore_cudagraph_manager(self, _config): + self.cudagraph_manager = object() + + config = TransformerConfig( + num_layers=1, + hidden_size=8, + num_attention_heads=2, + ffn_hidden_size=16, + cuda_graph_impl="local", + enable_hyper_connections=True, + num_residual_streams=2, + ) + inner = _LocalGraphInner(config) + wrapper = HyperConnectionHybridLayer(config, inner) + + assert hasattr(inner, "cudagraph_manager") + assert not hasattr(wrapper, "cudagraph_manager") + assert wrapper.offload_module_in_cuda_graph + + +def test_hybrid_wrapper_forwards_offload_stream_and_event_to_te_graph(monkeypatch): + import sys + from types import ModuleType, SimpleNamespace + + from megatron.core.models.hybrid.hybrid_block import HyperConnectionHybridLayer + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + FineGrainedActivationOffloadingInterface, + ) + + transformer_engine = ModuleType("transformer_engine") + transformer_engine.__path__ = [] + transformer_engine_pytorch = ModuleType("transformer_engine.pytorch") + transformer_engine.pytorch = transformer_engine_pytorch + monkeypatch.setitem(sys.modules, "transformer_engine", transformer_engine) + monkeypatch.setitem( + sys.modules, "transformer_engine.pytorch", transformer_engine_pytorch + ) + monkeypatch.setattr( + "megatron.core.transformer.transformer_layer.is_te_min_version", + lambda *_args, **_kwargs: True, + ) + + cuda_graph_stream = object() + cuda_graph_event = object() + monkeypatch.setattr( + FineGrainedActivationOffloadingInterface, + "cuda_graph_stream", + lambda: cuda_graph_stream, + ) + monkeypatch.setattr( + FineGrainedActivationOffloadingInterface, + "cuda_graph_event", + lambda: cuda_graph_event, + ) + + config = SimpleNamespace( + create_attention_mask_in_dataloader=False, + fine_grained_activation_offloading=True, + ) + inner = object.__new__(TransformerLayer) + torch.nn.Module.__init__(inner) + inner.config = config + inner.offload_module_in_cuda_graph = True + inner.current_microbatch = 7 + wrapper = object.__new__(HyperConnectionHybridLayer) + torch.nn.Module.__init__(wrapper) + wrapper.inner_layer = inner + wrapper.config = config + wrapper.current_microbatch = 1 + hidden_states = torch.ones(2, 1, 4) + + graph_args, graph_kwargs = wrapper._get_te_cuda_graph_replay_args( + hidden_states, attention_mask=None + ) + + assert len(graph_args) == 1 + assert graph_args[0] is hidden_states + assert graph_kwargs["is_first_microbatch"] is False + assert "attention_mask" not in graph_kwargs + assert graph_kwargs["cuda_graph_stream"] is cuda_graph_stream + assert graph_kwargs["cuda_graph_event"] is cuda_graph_event + assert inner.current_microbatch == 7 + + +def test_hybrid_wrapper_drops_runtime_packed_params_when_graph_has_no_contract( + monkeypatch, +): + from types import SimpleNamespace + + from megatron.core.models.hybrid.hybrid_block import HyperConnectionHybridLayer + from megatron.core.transformer.module import GraphableMegatronModule + + inner = object.__new__(TransformerLayer) + torch.nn.Module.__init__(inner) + inner.config = SimpleNamespace(delay_offload_until_cuda_graph=False) + wrapper = object.__new__(HyperConnectionHybridLayer) + torch.nn.Module.__init__(wrapper) + wrapper.inner_layer = inner + wrapper.config = SimpleNamespace(cuda_graph_modules=[]) + wrapper._te_cuda_graph_sample_kwarg_names = frozenset() + hidden_states = torch.ones(2, 1, 4) + replay_kwargs = {} + + def fake_graph_replay(_self, *args, **kwargs): + graph_hidden_states = ( + args[0] if args else kwargs.pop("hidden_states") + ) + assert graph_hidden_states is hidden_states + replay_kwargs.update(kwargs) + return (hidden_states,) + + monkeypatch.setattr( + GraphableMegatronModule, "_te_cuda_graph_replay", fake_graph_replay + ) + + replayed, context = wrapper._te_cuda_graph_replay( + hidden_states=hidden_states, + packed_seq_params=_make_packed_seq_params(), + ) + + assert replayed is hidden_states + assert context is None + assert "packed_seq_params" not in replay_kwargs + assert not has_packed_seq_params_cuda_graph_kwargs(replay_kwargs) + + +@pytest.mark.parametrize("layer_kind", ["attention", "moe", "hash_moe"]) +def test_hybrid_wrapper_replay_filters_to_captured_signature( + monkeypatch, layer_kind +): + from types import SimpleNamespace + + from megatron.core.models.hybrid.hybrid_block import HyperConnectionHybridLayer + from megatron.core.transformer.module import GraphableMegatronModule + + inner = object.__new__(TransformerLayer) + torch.nn.Module.__init__(inner) + inner.config = SimpleNamespace(delay_offload_until_cuda_graph=False) + inner.is_moe_layer = layer_kind != "attention" + if inner.is_moe_layer: + inner.mlp = SimpleNamespace( + router=SimpleNamespace(is_hash_layer=layer_kind == "hash_moe") + ) + + packed_seq_params = _make_packed_seq_params() + sample_kwarg_names = {"padding_mask"} + expected_flattened_kwargs = {} + if layer_kind == "attention": + expected_flattened_kwargs, static_metadata = ( + split_packed_seq_params_for_cuda_graph(packed_seq_params) + ) + inner._set_te_cuda_graph_packed_seq_params_static_metadata( + static_metadata, expected_flattened_kwargs + ) + sample_kwarg_names.update(expected_flattened_kwargs) + sample_kwarg_names.add("rotary_pos_emb") + elif layer_kind == "hash_moe": + sample_kwarg_names.add("input_ids") + + wrapper = object.__new__(HyperConnectionHybridLayer) + torch.nn.Module.__init__(wrapper) + wrapper.inner_layer = inner + wrapper.config = SimpleNamespace(cuda_graph_modules=[]) + wrapper._te_cuda_graph_sample_kwarg_names = frozenset(sample_kwarg_names) + + hidden_states = torch.ones(2, 1, 4) + padding_mask = torch.zeros(1, 2, dtype=torch.bool) + rotary_pos_emb = torch.ones(2, 1, 1, 4) + input_ids = torch.ones(1, 2, dtype=torch.long) + replay_kwargs = {} + + def fake_graph_replay(_self, *args, **kwargs): + graph_hidden_states = args[0] if args else kwargs.pop("hidden_states") + assert graph_hidden_states is hidden_states + replay_kwargs.update(kwargs) + return (hidden_states,) + + monkeypatch.setattr( + GraphableMegatronModule, "_te_cuda_graph_replay", fake_graph_replay + ) + + replayed, context = wrapper._te_cuda_graph_replay( + hidden_states=hidden_states, + attention_mask=torch.ones(1), + inference_context=object(), + rotary_pos_emb=rotary_pos_emb, + sequence_len_offset=torch.ones(1), + packed_seq_params=packed_seq_params, + padding_mask=padding_mask, + input_ids=input_ids, + ) + + assert replayed is hidden_states + assert context is None + assert set(replay_kwargs) == sample_kwarg_names + assert replay_kwargs["padding_mask"] is padding_mask + assert "packed_seq_params" not in replay_kwargs + assert has_packed_seq_params_cuda_graph_kwargs(replay_kwargs) == ( + layer_kind == "attention" + ) + if layer_kind == "attention": + assert replay_kwargs["rotary_pos_emb"] is rotary_pos_emb + assert set(expected_flattened_kwargs).issubset(replay_kwargs) + assert "input_ids" not in replay_kwargs + elif layer_kind == "hash_moe": + assert replay_kwargs["input_ids"] is input_ids + else: + assert "input_ids" not in replay_kwargs + + +def test_hybrid_wrapper_nonpacked_mamba_capture_and_replay(monkeypatch): + from types import SimpleNamespace + + from megatron.core.models.hybrid.hybrid_block import HyperConnectionHybridLayer + from megatron.core.ssm.mamba_layer import MambaLayer + from megatron.core.transformer.module import GraphableMegatronModule + + class _MambaStub(MambaLayer): + def __init__(self): + torch.nn.Module.__init__(self) + self.config = SimpleNamespace(cuda_graph_impl="none") + self.layer_number = 1 + + def forward(self, hidden_states, **_kwargs): + return hidden_states + + class _HyperConnectionStub: + def __call__(self, hidden_states, **_kwargs): + aggregated = hidden_states[..., :1] + return aggregated, torch.ones(1), torch.ones(1), hidden_states + + def fused_h_res_h_post_bda( + self, _h_res, residual, _h_post, _output_with_bias, **_kwargs + ): + return residual + + wrapper = object.__new__(HyperConnectionHybridLayer) + torch.nn.Module.__init__(wrapper) + wrapper.inner_layer = _MambaStub() + wrapper.hyper_connection = _HyperConnectionStub() + wrapper.config = SimpleNamespace( + fp32_residual_connection=False, params_dtype=None + ) + wrapper.training = True + hidden_states = torch.ones(2, 1, 2) + + captured = wrapper._te_cuda_graph_capture(hidden_states) + assert len(captured) == 1 + assert captured[0] is hidden_states + assert wrapper._te_cuda_graph_sample_kwarg_names == frozenset() + + replay_kwargs = {} + + def fake_graph_replay(_self, *_args, **kwargs): + replay_kwargs.update(kwargs) + return captured + + monkeypatch.setattr( + GraphableMegatronModule, "_te_cuda_graph_replay", fake_graph_replay + ) + replayed, context = wrapper._te_cuda_graph_replay( + hidden_states, + attention_mask=torch.ones(1), + inference_context=None, + rotary_pos_emb=torch.ones(1), + sequence_len_offset=torch.ones(1), + packed_seq_params=None, + padding_mask=torch.ones(1), + input_ids=torch.ones(1), + ) + + assert replayed is hidden_states + assert context is None + assert replay_kwargs == {} + + +def test_hybrid_wrapper_partial_moe_capture_and_replay_order(monkeypatch): + from types import SimpleNamespace + + from megatron.core.models.hybrid.hybrid_block import HyperConnectionHybridLayer + from megatron.core.transformer.module import GraphableMegatronModule + + aggregated = torch.tensor([1.0]) + h_res = torch.tensor([2.0]) + h_post = torch.tensor([3.0]) + residual = torch.tensor([4.0]) + router_hidden = torch.tensor([5.0]) + probs = torch.tensor([6.0]) + inner_residual = torch.tensor([7.0]) + expert_output = torch.tensor([8.0]) + final_output = torch.tensor([9.0]) + lifecycle_events = [] + + class _HyperConnectionStub: + def __call__(self, hidden_states, return_residual=False): + assert return_residual + return aggregated, h_res, h_post, residual + + def fused_h_res_h_post_bda( + self, + actual_h_res, + actual_residual, + actual_h_post, + output_with_bias, + **_kwargs, + ): + assert actual_h_res is h_res + assert actual_residual is residual + assert actual_h_post is h_post + assert output_with_bias == (expert_output, None) + lifecycle_events.append("bda") + return final_output + + class _OffloadInterfaceStub: + def enter_replay(self): + lifecycle_events.append("enter") + + def flush_delayed_groups(self): + lifecycle_events.append("flush") + + def exit_replay(self): + lifecycle_events.append("exit") + + class _InnerStub(TransformerLayer): + def __init__(self): + torch.nn.Module.__init__(self) + self.hidden_dropout = 0.0 + self.config = SimpleNamespace( + bias_dropout_fusion=False, + delay_offload_until_cuda_graph=True, + ) + self.off_interface = _OffloadInterfaceStub() + + def _rebuild_te_cuda_graph_packed_seq_params(self, kwargs): + return None + + def _flatten_te_cuda_graph_packed_seq_params(self, kwargs): + return None + + def _te_cuda_graph_capture(self, actual_aggregated, **_kwargs): + assert actual_aggregated is aggregated + return router_hidden, probs, inner_residual + + def resume_moe_experts_after_partial_cudagraph(self, outputs): + assert outputs == [router_hidden, probs, inner_residual] + lifecycle_events.append("experts") + return inner_residual, (expert_output, None) + + wrapper = object.__new__(HyperConnectionHybridLayer) + torch.nn.Module.__init__(wrapper) + wrapper.inner_layer = _InnerStub() + wrapper.hyper_connection = _HyperConnectionStub() + wrapper.config = SimpleNamespace( + fp32_residual_connection=False, + params_dtype=None, + bias_dropout_fusion=False, + ) + wrapper.training = True + monkeypatch.setattr( + HyperConnectionHybridLayer, + "_inner_is_partial_moe_capture", + lambda _self: True, + ) + + captured = wrapper._te_cuda_graph_capture(torch.tensor([0.0])) + assert captured == ( + router_hidden, + probs, + inner_residual, + h_post, + h_res, + residual, + ) + + monkeypatch.setattr( + GraphableMegatronModule, + "_te_cuda_graph_replay", + lambda _self, *_args, **_kwargs: ( + lifecycle_events.append("graph") or captured + ), + ) + replayed, context = wrapper._te_cuda_graph_replay(torch.tensor([0.0])) + assert replayed is final_output + assert context is None + assert lifecycle_events == ["enter", "graph", "flush", "experts", "bda", "exit"] + + +def test_transformer_layer_rejects_changed_packed_seq_params_tensor_fields(): + layer = _TransformerLayerCudaGraphStub() + packed_seq_params = _make_packed_seq_params() + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params) + layer._set_te_cuda_graph_packed_seq_params_static_metadata(static_metadata, tensor_kwargs) + packed_seq_params.cu_seqlens_q_padded = None + + with pytest.raises(AssertionError, match="Tensor fields"): + layer._flatten_te_cuda_graph_packed_seq_params({"packed_seq_params": packed_seq_params}) + + +def test_transformer_layer_rejects_replay_with_overlapping_flattened_kwargs(): + layer = _TransformerLayerCudaGraphStub() + packed_seq_params = _make_packed_seq_params() + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params) + layer._set_te_cuda_graph_packed_seq_params_static_metadata(static_metadata, tensor_kwargs) + existing_key = f"{CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX}cu_seqlens_q" + + with pytest.raises(AssertionError, match="overlap"): + layer._flatten_te_cuda_graph_packed_seq_params( + {existing_key: torch.IntTensor([0]), "packed_seq_params": packed_seq_params} + ) + + +def test_te_cuda_graph_sample_kwargs_include_flattened_packed_seq_params(): + layer = _TransformerLayerCudaGraphStub() + packed_seq_params = _make_packed_seq_params() + expected_tensor_kwargs, expected_static_metadata = split_packed_seq_params_for_cuda_graph( + packed_seq_params + ) + attention_mask = torch.zeros(1, 1, 16, 16, dtype=torch.bool) + sample_kwargs = {"attention_mask": attention_mask} + + _add_packed_seq_params_to_te_cuda_graph_sample_kwargs(layer, sample_kwargs, packed_seq_params) + + assert sample_kwargs["attention_mask"] is attention_mask + assert set(expected_tensor_kwargs).issubset(sample_kwargs) + for key, value in expected_tensor_kwargs.items(): + assert sample_kwargs[key] is value + assert layer._get_te_cuda_graph_packed_seq_params_static_metadata() == expected_static_metadata + assert layer._get_te_cuda_graph_packed_seq_params_tensor_kwarg_names() == tuple( + sorted(expected_tensor_kwargs) + ) + + +def test_te_cuda_graph_sample_kwargs_noop_without_packed_seq_params(): + layer = _TransformerLayerCudaGraphStub() + attention_mask = torch.zeros(1, 1, 16, 16, dtype=torch.bool) + sample_kwargs = {"attention_mask": attention_mask} + + _add_packed_seq_params_to_te_cuda_graph_sample_kwargs(layer, sample_kwargs, None) + + assert sample_kwargs == {"attention_mask": attention_mask} + assert layer._get_te_cuda_graph_packed_seq_params_static_metadata() is None + + +def test_te_cuda_graph_sample_kwargs_reject_overlapping_flattened_keys(): + layer = _TransformerLayerCudaGraphStub() + packed_seq_params = _make_packed_seq_params() + sample_kwargs = {f"{CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX}cu_seqlens_q": torch.IntTensor([0])} + + with pytest.raises(AssertionError, match="overlap"): + _add_packed_seq_params_to_te_cuda_graph_sample_kwargs( + layer, sample_kwargs, packed_seq_params + ) + + +def test_te_cuda_graph_partial_attn_only_flow(): + from megatron.core.transformer.enums import CudaGraphModule + + class _ConfigStub: + def __init__(self, cuda_graph_modules): + self.cuda_graph_modules = cuda_graph_modules + self.delay_offload_until_cuda_graph = False + + class _TestLayer(_TransformerLayerCudaGraphStub): + _te_cuda_graph_replay = TransformerLayer._te_cuda_graph_replay + + def __init__(self, cuda_graph_modules): + self.config = _ConfigStub(cuda_graph_modules) + self.attn_called = False + self.replay_impl_called = False + self.replay_impl_args = None + self.replay_impl_kwargs = None + self.replay_impl_context = None + + def _forward_attention(self, *args, **kwargs): + self.attn_called = True + return torch.ones(2, 1, 4) * 2.0, "attn_context" + + def _te_cuda_graph_replay_impl(self, args, kwargs, context): + self.replay_impl_called = True + self.replay_impl_args = args + self.replay_impl_kwargs = kwargs + self.replay_impl_context = context + return torch.ones(2, 1, 4) * 3.0 + + # Case 1: When CudaGraphModule.attn is captured + layer_attn = _TestLayer([CudaGraphModule.attn]) + packed_seq_params = _make_packed_seq_params() + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params) + layer_attn._set_te_cuda_graph_packed_seq_params_static_metadata(static_metadata, tensor_kwargs) + + kwargs = {"packed_seq_params": packed_seq_params, "hidden_states": torch.ones(2, 1, 4)} + layer_attn._te_cuda_graph_replay(**kwargs) + + assert not layer_attn.attn_called + assert layer_attn.replay_impl_called + assert layer_attn.replay_impl_context is None + assert "packed_seq_params" not in layer_attn.replay_impl_kwargs + assert f"{CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX}cu_seqlens_q" in layer_attn.replay_impl_kwargs + + # Case 2: When CudaGraphModule.attn is NOT captured (e.g. only mlp is captured) + layer_mlp = _TestLayer([CudaGraphModule.mlp]) + + kwargs = {"packed_seq_params": packed_seq_params, "hidden_states": torch.ones(2, 1, 4)} + layer_mlp._te_cuda_graph_replay(**kwargs) + + assert layer_mlp.attn_called + assert layer_mlp.replay_impl_called + assert layer_mlp.replay_impl_context == "attn_context" + assert len(layer_mlp.replay_impl_args) == 1 + assert torch.equal(layer_mlp.replay_impl_args[0], torch.ones(2, 1, 4) * 2.0) + assert layer_mlp.replay_impl_kwargs == {} + + +def test_seq_idx_determinism_across_replays(): + cu_seqlens = torch.IntTensor([0, 3, 7, 10]) + cu_seqlens_padded = torch.IntTensor([0, 4, 8, 12]) + + params1 = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=4, + max_seqlen_kv=4, + total_tokens=10, + ) + + params2 = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=4, + max_seqlen_kv=4, + total_tokens=10, + ) + + assert params1.seq_idx is not None + assert params2.seq_idx is not None + assert torch.equal(params1.seq_idx, params2.seq_idx) + assert params1.seq_idx.shape == params2.seq_idx.shape + assert params1.seq_idx.dtype == torch.int32 diff --git a/tests/unit_tests/transformer/test_thd_cuda_graph.py b/tests/unit_tests/transformer/test_thd_cuda_graph.py new file mode 100644 index 00000000000..9dc25141c06 --- /dev/null +++ b/tests/unit_tests/transformer/test_thd_cuda_graph.py @@ -0,0 +1,383 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Focused tests for static-shape THD Transformer Engine CUDA graphs.""" + +import pytest +import torch + +from megatron.core.packed_seq_params import ( + CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX, + PackedSeqParams, + get_thd_padding_kwargs, + pad_sequence_for_thd, +) +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.cuda_graphs import ( + TECudaGraphHelper, + _add_packed_seq_params_to_te_cuda_graph_sample_kwargs, +) +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.transformer.transformer_layer import TransformerLayer +from tests.unit_tests.test_utilities import Utils + + +def _make_cu(seqlens, device="cpu"): + lengths = torch.tensor(seqlens, dtype=torch.int32, device=device) + return torch.cat( + (torch.zeros(1, dtype=torch.int32, device=device), lengths.cumsum(0)) + ) + + +def _make_packed_seq_params(seqlens, device="cpu"): + cu_seqlens = _make_cu(seqlens, device=device) + return PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens.clone(), + cu_seqlens_q_padded=cu_seqlens.clone(), + cu_seqlens_kv_padded=cu_seqlens.clone(), + max_seqlen_q=max(seqlens), + max_seqlen_kv=max(seqlens), + local_cp_size=1, + pad_between_seqs=False, + ) + + +@pytest.mark.parametrize( + "cuda_graph_static,expected_max_num_seqs", [(False, None), (True, 32)] +) +def test_pad_to_max_resolves_static_cu_capacity( + cuda_graph_static, expected_max_num_seqs +): + alignment, target_len, max_num_seqs = get_thd_padding_kwargs( + pad_packed_seq_alignment="max", + max_seqlen_per_dp_cp_rank=8192, + thd_max_packed_sequences=32, + cuda_graph_static=cuda_graph_static, + ) + + assert alignment is None + assert target_len == 8192 + assert max_num_seqs == expected_max_num_seqs + + +def test_static_padding_appends_dummy_sequence_and_fixes_all_shapes(): + packed_seq_params = _make_packed_seq_params([100, 50, 30]) + packed_seq_params.cp_partition_mode = "contiguous" + total_tokens = 180 + tokens = torch.arange(total_tokens).reshape(1, -1) + + padded = pad_sequence_for_thd( + tokens, + tokens.clone(), + torch.ones(1, total_tokens), + torch.arange(total_tokens).reshape(1, -1), + packed_seq_params, + target_len=256, + max_num_seqs=8, + ) + padded_tokens, labels, loss_mask, position_ids, params, padding_mask = padded + + for tensor in (padded_tokens, labels, loss_mask, position_ids): + assert tensor.shape == (1, 256) + for cu_seqlens in ( + params.cu_seqlens_q, + params.cu_seqlens_kv, + params.cu_seqlens_q_padded, + params.cu_seqlens_kv_padded, + ): + assert cu_seqlens.shape == (9,) + assert cu_seqlens.tolist() == [0, 100, 150, 180, 256, 256, 256, 256, 256] + assert params.max_seqlen_q == 256 + assert params.max_seqlen_kv == 256 + assert params.pad_between_seqs is False + assert params.local_cp_size == packed_seq_params.local_cp_size + assert params.cp_partition_mode == "contiguous" + assert padding_mask.shape == (1, 256) + assert not padding_mask[0, :total_tokens].any() + assert padding_mask[0, total_tokens:].all() + + +def test_alignment_padding_preserves_metadata_without_dummy_sequence(): + packed_seq_params = _make_packed_seq_params([50, 30]) + original_cu_seqlens = packed_seq_params.cu_seqlens_q.clone() + + padded_tokens, _, _, _, params, padding_mask = pad_sequence_for_thd( + torch.ones(1, 80), + None, + None, + None, + packed_seq_params, + alignment=64, + pad_by_appending_dummy_seq=False, + ) + + assert padded_tokens.shape == (1, 128) + assert torch.equal(params.cu_seqlens_q, original_cu_seqlens) + assert params.max_seqlen_q == 50 + assert params.pad_between_seqs is False + assert not padding_mask[0, :80].any() + assert padding_mask[0, 80:].all() + + +def test_padding_merges_existing_mask_with_tail(): + packed_seq_params = _make_packed_seq_params([4, 4]) + existing_mask = torch.tensor( + [[False, False, False, True, False, False, True, True]] + ) + + *_, padding_mask = pad_sequence_for_thd( + torch.ones(1, 8), + None, + None, + None, + packed_seq_params, + target_len=10, + max_num_seqs=4, + padding_mask=existing_mask, + ) + + assert padding_mask.tolist() == [ + [False, False, False, True, False, False, True, True, True, True] + ] + + +def test_metadata_only_padding_uses_explicit_cp_geometry(): + packed_seq_params = _make_packed_seq_params([15]) + packed_seq_params.cu_seqlens_q_padded = torch.tensor( + [0, 16], dtype=torch.int32 + ) + packed_seq_params.cu_seqlens_kv_padded = ( + packed_seq_params.cu_seqlens_q_padded.clone() + ) + existing_mask = torch.tensor([[False, False, False, True]]) + + *_, params, padding_mask = pad_sequence_for_thd( + None, + None, + None, + None, + packed_seq_params, + target_len=4, + padding_mask=existing_mask, + cp_size=4, + cp_rank=3, + ) + + assert torch.equal(padding_mask, existing_mask) + assert params.cu_seqlens_q.tolist() == [0, 15, 16] + assert params.cu_seqlens_q_padded.tolist() == [0, 16, 16] + + +def test_dynamic_slot_liveness_for_pp_and_vpp_orders(): + pp_order = [1, 1, -1, 1, -1, 1, -1, -1] + vpp_order = [ + 1, + 1, + 1, + 2, + 2, + 2, + -2, + 1, + -2, + 1, + -2, + 2, + -1, + 2, + -1, + -1, + -2, + -2, + -1, + -1, + ] + + assert TECudaGraphHelper._get_required_num_microbatch_slots_from_order( + pp_order, 1 + ) == 2 + assert TECudaGraphHelper._get_required_num_microbatch_slots_from_order( + vpp_order, 2 + ) == 5 + + +def test_dp_balanced_capture_upper_bound_accounts_for_cp_and_vpp(): + assert ( + TECudaGraphHelper._get_dp_balanced_thd_max_num_microbatches( + global_batch_size=64, + dp_size=1, + cp_size=2, + max_seqlen_per_dp_cp_rank=4096, + max_sequence_length=4096, + max_num_seqs=8, + ) + == 32 + ) + assert ( + TECudaGraphHelper._get_dp_balanced_thd_max_num_microbatches( + global_batch_size=18, + dp_size=1, + cp_size=1, + max_seqlen_per_dp_cp_rank=4096, + max_sequence_length=2048, + microbatch_group_size_per_vp_stage=8, + max_num_seqs=8, + ) + == 16 + ) + + +def test_thd_te_graph_rejects_moe_tp_sp_token_count_mismatch(): + with pytest.raises( + ValueError, + match="tensor_parallel_size > 1 with sequence_parallel", + ): + TransformerConfig( + num_layers=1, + hidden_size=128, + num_attention_heads=4, + ffn_hidden_size=256, + num_moe_experts=8, + tensor_model_parallel_size=2, + sequence_parallel=True, + max_seqlen_per_dp_cp_rank=128, + sequence_packing_scheduler="dp_balanced", + pad_packed_seq_alignment="max", + cuda_graph_impl="transformer_engine", + ) + + +def test_local_graph_does_not_require_static_thd_padding_contract(): + config = TransformerConfig( + num_layers=1, + hidden_size=128, + num_attention_heads=4, + ffn_hidden_size=256, + max_seqlen_per_dp_cp_rank=128, + sequence_packing_scheduler="dp_balanced", + cuda_graph_impl="local", + ) + + assert config.pad_packed_seq_alignment is None + + +@pytest.mark.parametrize( + "cuda_graph_kwargs", + [ + {"cuda_graph_impl": "transformer_engine"}, + {"external_cuda_graph": True}, + ], +) +def test_te_cuda_graph_rejects_mhc_selective_recompute(cuda_graph_kwargs): + with pytest.raises( + NotImplementedError, + match="'mhc' in recompute_modules is not supported", + ): + TransformerConfig( + num_layers=1, + hidden_size=128, + num_attention_heads=4, + ffn_hidden_size=256, + enable_hyper_connections=True, + num_residual_streams=2, + recompute_granularity="selective", + recompute_modules=["mhc"], + **cuda_graph_kwargs, + ) + + +@pytest.mark.parametrize( + "unsupported_recompute_kwargs", + [ + {"cuda_graph_impl": "transformer_engine"}, + {"external_cuda_graph": True}, + { + "cuda_graph_impl": "local", + "cuda_graph_modules": ["moe_router"], + "fine_grained_activation_offloading": True, + "offload_modules": ["expert_fc1"], + "num_moe_experts": 4, + }, + ], +) +def test_mhc_recompute_warning_is_suppressed_when_unsupported( + recwarn, unsupported_recompute_kwargs +): + TransformerConfig( + num_layers=1, + hidden_size=128, + num_attention_heads=4, + ffn_hidden_size=256, + enable_hyper_connections=True, + num_residual_streams=2, + recompute_granularity=None, + recompute_modules=[], + **unsupported_recompute_kwargs, + ) + + assert not any( + "Consider adding 'mhc'" in str(warning.message) for warning in recwarn + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_transformer_layer_static_thd_inputs_use_prefixed_contract(): + from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec + + Utils.initialize_model_parallel(tensor_model_parallel_size=1) + try: + config = TransformerConfig( + num_layers=1, + hidden_size=256, + num_attention_heads=4, + ffn_hidden_size=1024, + max_seqlen_per_dp_cp_rank=128, + sequence_packing_scheduler="dp_balanced", + pad_packed_seq_alignment="max", + thd_max_packed_sequences=8, + cuda_graph_impl="transformer_engine", + cp_partition_mode="contiguous", + bf16=True, + ) + model_parallel_cuda_manual_seed(42) + attention_layer_spec = hybrid_stack_spec.submodules.attention_layer + layer = ( + TransformerLayer( + config, + attention_layer_spec.submodules, + layer_number=1, + ) + .cuda() + .bfloat16() + ) + + static_inputs = layer.get_layer_static_inputs( + seq_length=128, micro_batch_size=4 + ) + packed_seq_params = static_inputs.pop("packed_seq_params") + assert packed_seq_params.cp_partition_mode == "contiguous" + _add_packed_seq_params_to_te_cuda_graph_sample_kwargs( + layer, static_inputs, packed_seq_params + ) + assert ( + layer._get_te_cuda_graph_packed_seq_params_static_metadata()[ + "cp_partition_mode" + ] + == "contiguous" + ) + + assert static_inputs["hidden_states"].shape == (128, 1, 256) + assert static_inputs["hidden_states"].dtype == torch.bfloat16 + assert static_inputs["padding_mask"].shape == (1, 128) + assert not static_inputs["padding_mask"].any() + assert not any(key.startswith("cu_seqlens_") for key in static_inputs) + flattened_keys = { + key + for key in static_inputs + if key.startswith(CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX) + } + assert len(flattened_keys) == 4 + assert all(static_inputs[key].shape == (9,) for key in flattened_keys) + finally: + Utils.destroy_model_parallel() diff --git a/tests/unit_tests/transformer/test_transformer_layer.py b/tests/unit_tests/transformer/test_transformer_layer.py index 93650cf13b0..d400655673f 100644 --- a/tests/unit_tests/transformer/test_transformer_layer.py +++ b/tests/unit_tests/transformer/test_transformer_layer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import gc @@ -16,13 +16,15 @@ ) from megatron.core.tensor_parallel.random import ( HAVE_TE, + CheckpointWithoutOutputManager, initialize_rng_tracker, model_parallel_cuda_manual_seed, ) from megatron.core.transformer.cuda_graphs import CudaGraphManager, _CudagraphGlobalRecord -from megatron.core.transformer.enums import InferenceCudaGraphScope +from megatron.core.transformer.enums import AttnMaskType, CudaGraphModule, InferenceCudaGraphScope from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.transformer_layer import ( + HyperConnectionTransformerLayer, TransformerLayer, get_transformer_layer_offset, ) @@ -30,6 +32,44 @@ from tests.unit_tests.test_utilities import Utils +def _make_mhc_layer_spec(**kwargs): + """Build a layer spec with HyperConnectionModule submodules. + + The ``enable_hyper_connection`` kwarg on ``gpt_layer_specs`` is added by + the GPT-wiring follow-up split, so this helper patches the mHC submodules + directly to keep the unit tests self-contained for this split. + """ + from megatron.core.transformer.hyper_connection import HyperConnectionModule + + layer_spec = get_gpt_layer_with_transformer_engine_spec(**kwargs) + layer_spec.module = HyperConnectionTransformerLayer + layer_spec.submodules.self_attention_hyper_connection = HyperConnectionModule + layer_spec.submodules.mlp_hyper_connection = HyperConnectionModule + return layer_spec + + +def _make_mhc_config(hidden_size=64, num_streams=4, **extra): + """Build a TransformerConfig with common MHC defaults. + + Any default can be overridden via **extra + (e.g. ``_make_mhc_config(num_layers=8, recompute_modules=["core_attn", "mhc"])``). + """ + base = dict( + num_layers=2, + hidden_size=hidden_size, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=num_streams, + mhc_sinkhorn_iterations=5, + mhc_init_gating_factor=0.01, + hidden_dropout=0.0, + attention_dropout=0.0, + ) + base.update(extra) + return TransformerConfig(**base) + + class TestParallelTransformerLayer: def setup_method(self, method): @@ -73,6 +113,97 @@ def test_gpu_forward(self): assert hidden_states.shape[1] == micro_batch_size assert hidden_states.shape[2] == config.hidden_size + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + @pytest.mark.skipif( + not (HAVE_TE and is_te_min_version("1.10.0")), + reason="TE CUDA graph kwargs require TransformerEngine >= 1.10", + ) + def test_te_cuda_graph_omits_absent_attention_mask(self): + config = TransformerConfig( + num_layers=2, + hidden_size=12, + num_attention_heads=4, + use_cpu_initialization=True, + cuda_graph_impl="transformer_engine", + cuda_graph_modules=[CudaGraphModule.attn], + create_attention_mask_in_dataloader=False, + ) + layer = TransformerLayer( + config, get_gpt_layer_with_transformer_engine_submodules() + ) + + static_inputs = layer.get_layer_static_inputs( + seq_length=32, micro_batch_size=1 + ) + assert "attention_mask" not in static_inputs + + hidden_states = torch.ones((32, 1, config.hidden_size), device="cuda") + _, cudagraph_kwargs = layer._get_te_cuda_graph_replay_args( + hidden_states, attention_mask=None + ) + assert "attention_mask" not in cudagraph_kwargs + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + @pytest.mark.skipif( + not (HAVE_TE and is_te_min_version("1.10.0")), + reason="TE CUDA graph kwargs require TransformerEngine >= 1.10", + ) + def test_te_cuda_graph_keeps_configured_attention_mask(self): + config = TransformerConfig( + num_layers=2, + hidden_size=12, + num_attention_heads=4, + use_cpu_initialization=True, + cuda_graph_impl="transformer_engine", + cuda_graph_modules=[CudaGraphModule.attn], + create_attention_mask_in_dataloader=True, + ) + layer = TransformerLayer( + config, get_gpt_layer_with_transformer_engine_submodules() + ) + + static_inputs = layer.get_layer_static_inputs( + seq_length=32, micro_batch_size=2 + ) + assert static_inputs["attention_mask"].shape == (2, 1, 32, 32) + + hidden_states = torch.ones((32, 2, config.hidden_size), device="cuda") + _, cudagraph_kwargs = layer._get_te_cuda_graph_replay_args( + hidden_states, attention_mask=None + ) + assert cudagraph_kwargs["attention_mask"].shape == (2, 1, 32, 32) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + @pytest.mark.skipif( + not (HAVE_TE and is_te_min_version("1.10.0")), + reason="TE CUDA graph kwargs require TransformerEngine >= 1.10", + ) + def test_te_cuda_graph_warns_for_omitted_padding_mask(self, monkeypatch): + config = TransformerConfig( + num_layers=2, + hidden_size=12, + num_attention_heads=4, + use_cpu_initialization=True, + cuda_graph_impl="transformer_engine", + cuda_graph_modules=[CudaGraphModule.attn], + create_attention_mask_in_dataloader=False, + ) + layer = TransformerLayer( + config, get_gpt_layer_with_transformer_engine_submodules() + ) + layer.self_attention.attn_mask_type = AttnMaskType.padding + + warnings = [] + monkeypatch.setattr( + "megatron.core.transformer.transformer_layer.log_single_rank", + lambda _logger, _level, message: warnings.append(message), + ) + static_inputs = layer.get_layer_static_inputs( + seq_length=32, micro_batch_size=1 + ) + assert "attention_mask" not in static_inputs + assert any("attn_mask_type=padding" in message for message in warnings) + def test_chunked_mlp(self): with torch.no_grad(): num_layers = 2 @@ -418,3 +549,767 @@ def test_deprecated_full_iteration_inference_scope_string_matches_new_granularit assert block.config.cuda_graph_modules == [] assert _no_layers_have_manager(block) _reset_cudagraph_state() + + +class TestTransformerLayerWithHyperConnectionRecompute: + """Test TransformerLayer with HyperConnection and MHC block recomputation.""" + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def _create_layer_with_hyper_connection( + self, hidden_size=64, num_streams=4, layer_number=1, **extra + ): + """Create a HyperConnectionTransformerLayer with hyper connection enabled.""" + config = _make_mhc_config( + hidden_size=hidden_size, + num_streams=num_streams, + recompute_modules=["core_attn", "mhc"], + recompute_granularity='selective', + **extra, + ) + layer_spec = _make_mhc_layer_spec() + layer = HyperConnectionTransformerLayer( + config, layer_spec.submodules, layer_number=layer_number + ) + layer.cuda() + return layer, config + + def test_forward_with_hyper_connection_recompute(self): + """ + Test that TransformerLayer forward works correctly with HyperConnection + and MHC block recomputation enabled. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + layer, config = self._create_layer_with_hyper_connection(hidden_size, num_streams) + layer.train() # Enable training mode for recomputation + + # Input shape: [seq_len, batch_size, n * hidden_size] for hyper connections + n_channels = num_streams * hidden_size + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + # Create manager for MHC block recomputation + manager = CheckpointWithoutOutputManager() + + # Forward pass with recompute manager + manager.is_last_layer_in_recompute_block = True + output, context = layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + mhc_recompute_manager=manager, + ) + + # Verify output shape + assert output.shape == ( + seq_len, + batch_size, + n_channels, + ), f"Expected output shape {(seq_len, batch_size, n_channels)}, got {output.shape}" + + # Register unified recompute hook at block boundary. + manager.discard_all_outputs_and_register_unified_recompute(output) + + # Backward pass should work without error + loss = output.sum() + loss.backward() + + # Verify gradients exist + assert hidden_states.grad is not None, "Gradients should be computed for hidden_states" + assert hidden_states.grad.shape == hidden_states.shape + + def test_intermediate_layer_with_recompute(self): + """ + Test TransformerLayer as an intermediate layer (not last in block). + In this case, MLP BDA should also be checkpointed. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + layer, config = self._create_layer_with_hyper_connection(hidden_size, num_streams) + layer.train() + + n_channels = num_streams * hidden_size + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + manager = CheckpointWithoutOutputManager() + + # Forward pass - NOT the last layer in block + manager.is_last_layer_in_recompute_block = False + output, context = layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + mhc_recompute_manager=manager, + ) + + # Verify output shape + assert output.shape == (seq_len, batch_size, n_channels) + + # Backward pass should work + loss = output.sum() + # For intermediate layers, we need to pass output to next layer + # Here we just register the recompute hook on output for testing + manager.discard_all_outputs_and_register_unified_recompute(loss) + + loss.backward() + + assert hidden_states.grad is not None + assert hidden_states.grad.shape == hidden_states.shape + + def test_multiple_layers_chain_with_recompute(self): + """ + Test multiple TransformerLayers chained together with a single + CheckpointWithoutOutputManager, simulating TransformerBlock behavior. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + num_layers = 3 + + layers = [ + self._create_layer_with_hyper_connection( + hidden_size, num_streams, layer_number=i + 1, num_layers=num_layers + )[0] + for i in range(num_layers) + ] + + for layer in layers: + layer.train() + + n_channels = num_streams * hidden_size + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + # Single manager for all layers (like TransformerBlock) + manager = CheckpointWithoutOutputManager() + + # Forward through all layers + h = hidden_states + for i, layer in enumerate(layers): + is_last = i == num_layers - 1 + manager.is_last_layer_in_recompute_block = is_last + h, _ = layer( + hidden_states=h, attention_mask=attention_mask, mhc_recompute_manager=manager + ) + if is_last: + manager.discard_all_outputs_and_register_unified_recompute(h) + + # Backward pass + loss = h.sum() + loss.backward() + + # Verify gradients + assert hidden_states.grad is not None + assert hidden_states.grad.shape == hidden_states.shape + # Check that gradient is non-trivial (not all zeros) + assert hidden_states.grad.abs().sum() > 0 + + +class TestMHCRecomputeMemorySaving: + """Verify that 'mhc' in recompute_modules actually reduces peak GPU memory.""" + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + @staticmethod + def _run_forward_backward( + num_layers, + hidden_size, + num_streams, + seq_len, + batch_size, + use_recompute, + recompute_block_size=2, + ): + """Run a full forward + backward pass and return (peak memory, output grad). + + When use_recompute=True, a new CheckpointWithoutOutputManager is created every + `recompute_block_size` layers, mirroring TransformerBlock's + _build_mhc_recompute_layer_plan logic. + """ + config = _make_mhc_config( + hidden_size=hidden_size, + num_streams=num_streams, + num_layers=num_layers, + recompute_modules=["core_attn", "mhc"] if use_recompute else None, + recompute_granularity='selective' if use_recompute else None, + ) + layer_spec = _make_mhc_layer_spec() + layers = [ + HyperConnectionTransformerLayer( + config, layer_spec.submodules, layer_number=i + 1 + ).cuda() + for i in range(num_layers) + ] + for layer in layers: + layer.train() + + n_channels = num_streams * hidden_size + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + torch.cuda.reset_peak_memory_stats() + torch.cuda.synchronize() + + manager = CheckpointWithoutOutputManager() if use_recompute else None + + h = hidden_states + for i, layer in enumerate(layers): + is_last_in_block = (i == num_layers - 1) or ((i + 1) % recompute_block_size == 0) + kwargs = dict(hidden_states=h, attention_mask=attention_mask) + if manager is not None: + manager.is_last_layer_in_recompute_block = is_last_in_block + kwargs['mhc_recompute_manager'] = manager + h, _ = layer(**kwargs) + if manager is not None and is_last_in_block: + manager.discard_all_outputs_and_register_unified_recompute(h) + if i < num_layers - 1: + manager = CheckpointWithoutOutputManager() + + loss = h.sum() + loss.backward() + torch.cuda.synchronize() + + peak_mem = torch.cuda.max_memory_allocated() + grad = hidden_states.grad.clone() + + del layers, hidden_states, h, loss, manager + torch.cuda.empty_cache() + + return peak_mem, grad + + def test_recompute_reduces_peak_memory(self): + """Peak memory with recompute (block_size=2) should be lower than without.""" + num_layers = 8 + hidden_size = 128 + num_streams = 4 + seq_len = 64 + batch_size = 4 + + peak_no_recompute, _ = self._run_forward_backward( + num_layers, hidden_size, num_streams, seq_len, batch_size, use_recompute=False + ) + peak_recompute, _ = self._run_forward_backward( + num_layers, + hidden_size, + num_streams, + seq_len, + batch_size, + use_recompute=True, + recompute_block_size=2, + ) + + saving_pct = (peak_no_recompute - peak_recompute) / peak_no_recompute * 100 + + assert peak_recompute < peak_no_recompute, ( + f"Recompute should reduce peak memory, but got " + f"no_recompute={peak_no_recompute / 1e6:.1f}MB vs " + f"recompute={peak_recompute / 1e6:.1f}MB " + f"(saving={saving_pct:.1f}%)" + ) + + +class TestMHCWithCudaGraph: + """Test HyperConnectionTransformerLayer compatibility with CUDA graphs. + + CUDA graph capture requires static computation graphs and fixed tensor shapes. + These tests verify that the mHC layer properly supports the CUDA graph interface + defined in GraphableMegatronModule and TransformerLayer. + """ + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123, use_cudagraphable_rng=True, force_reset_rng=True) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def _create_mhc_layer(self, hidden_size=64, num_streams=4, **extra_config): + config = _make_mhc_config(hidden_size=hidden_size, num_streams=num_streams, **extra_config) + layer_spec = _make_mhc_layer_spec() + layer = HyperConnectionTransformerLayer(config, layer_spec.submodules) + layer.cuda() + return layer, config + + def test_get_layer_static_inputs_shape_for_mhc(self): + """get_layer_static_inputs must return [s, b, n*C] for mHC layers. + + CUDA graph capture creates static buffers whose shapes are determined by + this method. If the shape is [s, b, C] instead of [s, b, n*C], the graph + capture will produce a shape mismatch at the first hyper connection module. + """ + layer, config = self._create_mhc_layer() + seq_length = 32 + micro_batch_size = 2 + + static_inputs = layer.get_layer_static_inputs(seq_length, micro_batch_size) + hidden_states = static_inputs["hidden_states"] + + expected_hidden_dim = config.num_residual_streams * config.hidden_size + assert hidden_states.shape[-1] == expected_hidden_dim, ( + f"get_layer_static_inputs returns hidden dim {hidden_states.shape[-1]} " + f"but mHC expects {expected_hidden_dim} (n={config.num_residual_streams} * " + f"C={config.hidden_size}). " + f"HyperConnectionTransformerLayer must override get_layer_static_inputs." + ) + + def test_submodules_under_cudagraphs_includes_hyper_connection(self): + """_get_submodules_under_cudagraphs must include hyper connection modules. + + CUDA graph manual hooks are set up for parameters of submodules returned + by this method. Missing hyper connection modules means their parameters + (mapping_proj, alpha_*, bias) will not get proper pre-forward hooks during + graph replay, leading to stale parameter values. + """ + layer, config = self._create_mhc_layer() + + submodules = layer._get_submodules_under_cudagraphs() + + hc_modules_found = any( + hasattr(m, 'mapping_proj') for submod in submodules for m in submod.modules() + ) + assert hc_modules_found, ( + "_get_submodules_under_cudagraphs does not include HyperConnectionModule. " + "Parameters like mapping_proj, alpha_pre/post/res will not be updated " + "during CUDA graph replay." + ) + + def test_forward_through_te_cuda_graph_capture_path(self): + """_te_cuda_graph_capture must produce correct output shapes for mHC. + + TE CUDA graph capture calls _te_cuda_graph_capture() during warmup. + For mHC layers, the input must be n-stream [s, b, n*C] and output must + also be [s, b, n*C]. + """ + layer, config = self._create_mhc_layer() + layer.eval() + + seq_len = 8 + batch_size = 2 + n_channels = config.num_residual_streams * config.hidden_size + + hidden_states = torch.randn(seq_len, batch_size, n_channels, device='cuda') + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + with torch.no_grad(): + outputs = layer._te_cuda_graph_capture( + hidden_states=hidden_states, attention_mask=attention_mask + ) + + if isinstance(outputs, tuple): + output = outputs[0] + else: + output = outputs + + assert output.shape == (seq_len, batch_size, n_channels), ( + f"_te_cuda_graph_capture output shape {output.shape} != " + f"expected {(seq_len, batch_size, n_channels)}" + ) + + def test_cuda_graph_fwd_bwd_with_hyper_connection(self): + """End-to-end CUDA graph capture and replay for forward+backward with mHC. + + Captures both the forward and backward pass of HyperConnectionTransformerLayer + into a torch.cuda.CUDAGraph and replays it with fresh input data, verifying + that the computation graph is fully static (capturable) and produces correct + output shapes and non-trivial gradients. + """ + layer, config = self._create_mhc_layer() + layer.train() + + seq_len = 8 + batch_size = 2 + n_channels = config.num_residual_streams * config.hidden_size + + static_input = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + # Warmup on side stream to trigger lazy allocations + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + out, _ = layer(hidden_states=static_input, attention_mask=attention_mask) + out.sum().backward() + torch.cuda.current_stream().wait_stream(s) + + # Set .grad to None so backward allocates fresh gradient tensors in the + # graph's private memory pool during capture. + layer.zero_grad(set_to_none=True) + static_input.grad = None + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + output, _ = layer(hidden_states=static_input, attention_mask=attention_mask) + output.sum().backward() + + # Replay with new input data. + # Use no_grad because backward inside the captured graph already + # bumped the autograd version counter on static_input, making + # in-place copy_ illegal without disabling grad tracking. + with torch.no_grad(): + static_input.copy_(torch.randn_like(static_input)) + g.replay() + + assert output.shape == ( + seq_len, + batch_size, + n_channels, + ), f"Output shape {output.shape} != expected {(seq_len, batch_size, n_channels)}" + assert ( + static_input.grad is not None + ), "Gradients should be computed for static_input after graph replay" + assert static_input.grad.shape == static_input.shape + assert static_input.grad.abs().sum() > 0, "Gradients should be non-trivial" + + # Verify numerical consistency: graph replay should match eager execution + # with the same input and weights. + test_data = torch.randn(seq_len, batch_size, n_channels, device='cuda') + + with torch.no_grad(): + static_input.copy_(test_data) + g.replay() + graph_out = output.detach().clone() + graph_grad = static_input.grad.detach().clone() + + eager_input = test_data.clone().requires_grad_(True) + eager_output, _ = layer(hidden_states=eager_input, attention_mask=attention_mask) + eager_output.sum().backward() + + assert torch.allclose(graph_out, eager_output.detach(), atol=1e-5), ( + f"Graph vs eager output mismatch: " + f"max diff = {(graph_out - eager_output.detach()).abs().max().item()}" + ) + assert torch.allclose(graph_grad, eager_input.grad, atol=1e-5), ( + f"Graph vs eager gradient mismatch: " + f"max diff = {(graph_grad - eager_input.grad).abs().max().item()}" + ) + + def test_cuda_graph_fwd_bwd_with_hyper_connection_and_recompute(self): + """CUDA graph capture+replay for fwd+bwd with mHC and CheckpointWithoutOutputManager. + + When a CheckpointWithoutOutputManager is used, additional CheckpointWithoutOutput + objects are created for layernorm and hyper-connection operations. The + manager discards intermediate activations during forward (storage.resize_(0)) + and recomputes them during backward via a unified gradient hook. + This test verifies the full capture+replay still works correctly. + """ + layer, config = self._create_mhc_layer() + layer.train() + + seq_len = 8 + batch_size = 2 + n_channels = config.num_residual_streams * config.hidden_size + + static_input = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + # Warmup on side stream; fresh manager per iteration to avoid stale state. + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + mgr = CheckpointWithoutOutputManager() + mgr.is_last_layer_in_recompute_block = True + out, _ = layer( + hidden_states=static_input, + attention_mask=attention_mask, + mhc_recompute_manager=mgr, + ) + mgr.discard_all_outputs_and_register_unified_recompute(out) + out.sum().backward() + torch.cuda.current_stream().wait_stream(s) + + layer.zero_grad(set_to_none=True) + static_input.grad = None + + capture_mgr = CheckpointWithoutOutputManager() + capture_mgr.is_last_layer_in_recompute_block = True + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + output, _ = layer( + hidden_states=static_input, + attention_mask=attention_mask, + mhc_recompute_manager=capture_mgr, + ) + capture_mgr.discard_all_outputs_and_register_unified_recompute(output) + output.sum().backward() + + # Replay with new input data. + with torch.no_grad(): + static_input.copy_(torch.randn_like(static_input)) + g.replay() + + assert output.shape == ( + seq_len, + batch_size, + n_channels, + ), f"Output shape {output.shape} != expected {(seq_len, batch_size, n_channels)}" + assert ( + static_input.grad is not None + ), "Gradients should be computed for static_input after graph replay" + assert static_input.grad.shape == static_input.shape + assert static_input.grad.abs().sum() > 0, "Gradients should be non-trivial" + + # Numerical consistency: graph replay vs eager with the same input. + test_data = torch.randn(seq_len, batch_size, n_channels, device='cuda') + + with torch.no_grad(): + static_input.copy_(test_data) + g.replay() + graph_out = output.detach().clone() + graph_grad = static_input.grad.detach().clone() + + eager_mgr = CheckpointWithoutOutputManager() + eager_mgr.is_last_layer_in_recompute_block = True + eager_input = test_data.clone().requires_grad_(True) + eager_output, _ = layer( + hidden_states=eager_input, + attention_mask=attention_mask, + mhc_recompute_manager=eager_mgr, + ) + eager_mgr.discard_all_outputs_and_register_unified_recompute(eager_output) + eager_output.sum().backward() + + assert torch.allclose(graph_out, eager_output.detach(), atol=1e-5), ( + f"Graph vs eager output mismatch: " + f"max diff = {(graph_out - eager_output.detach()).abs().max().item()}" + ) + assert torch.allclose(graph_grad, eager_input.grad, atol=1e-5), ( + f"Graph vs eager gradient mismatch: " + f"max diff = {(graph_grad - eager_input.grad).abs().max().item()}" + ) + + def test_mcore_cudagraph_manager_with_mhc_recompute_manager(self): + """MCore CudaGraphManager must not crash on mhc_recompute_manager kwarg. + + When cuda_graph_impl="local" is set, HyperConnectionTransformerLayer.__call__ + runs first and pops mhc_recompute_manager off kwargs before + super().__call__ → MegatronModule.__call__ → CudaGraphManager.__call__, + which iterates over all kwargs to check supported types. + CheckpointWithoutOutputManager (used by mhc_recompute_manager) is not a + CUDA-graph-supported type. + + This test verifies that mhc_recompute_manager is properly extracted + from kwargs before the CudaGraphManager sees them, preventing the + AssertionError that would otherwise occur. + """ + layer, config = self._create_mhc_layer(cuda_graph_impl="local", cuda_graph_scope="attn") + layer.train() + + assert hasattr( + layer, 'cudagraph_manager' + ), "Layer should have cudagraph_manager with cuda_graph_impl='local'" + + seq_len = 8 + batch_size = 2 + n_channels = config.num_residual_streams * config.hidden_size + + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + mgr = CheckpointWithoutOutputManager() + mgr.is_last_layer_in_recompute_block = True + + output, context = layer( + hidden_states=hidden_states, attention_mask=attention_mask, mhc_recompute_manager=mgr + ) + + assert output.shape == (seq_len, batch_size, n_channels) + + def test_mcore_cudagraph_manager_without_mhc_recompute_manager(self): + """MCore CudaGraphManager path works when mhc_recompute_manager is None.""" + layer, config = self._create_mhc_layer(cuda_graph_impl="local", cuda_graph_scope="attn") + layer.train() + + seq_len = 8 + batch_size = 2 + n_channels = config.num_residual_streams * config.hidden_size + + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + output, context = layer(hidden_states=hidden_states, attention_mask=attention_mask) + + assert output.shape == (seq_len, batch_size, n_channels) + + +class TestMHCWithOffloading: + """Test HyperConnectionTransformerLayer with fine-grained activation offloading. + + Fine-grained activation offloading transfers specific activations (e.g., layernorm + inputs) to CPU during forward and reloads them during backward. These tests verify + that the mHC layer's multi-stream architecture works correctly with offloading. + """ + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def _create_mhc_layer_with_offloading( + self, hidden_size=64, num_streams=4, offload_modules=None + ): + if offload_modules is None: + offload_modules = ["attn_norm", "mlp_norm"] + + config = _make_mhc_config( + hidden_size=hidden_size, + num_streams=num_streams, + fine_grained_activation_offloading=True, + offload_modules=offload_modules, + ) + layer_spec = _make_mhc_layer_spec() + layer = HyperConnectionTransformerLayer(config, layer_spec.submodules) + layer.cuda() + return layer, config + + def test_forward_backward_with_offloading(self): + """Forward+backward should work with activation offloading enabled. + + This exercises the off_interface context manager around layernorms in + the mHC forward path, including the group_commit that commits the + offloading group for the aggregated 1-stream layernorm input. + """ + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + PipelineOffloadManager, + ) + + layer, config = self._create_mhc_layer_with_offloading() + layer.train() + + seq_len = 8 + batch_size = 2 + n_channels = config.num_residual_streams * config.hidden_size + + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + mgr = PipelineOffloadManager.get_instance() + mgr.init_model_chunk_offload_handler( + pp_rank=0, vp_size=1, vp_stage=0, min_offloaded_tensor_size=0 + ) + + output, context = layer(hidden_states=hidden_states, attention_mask=attention_mask) + + assert output.shape == ( + seq_len, + batch_size, + n_channels, + ), f"Output shape {output.shape} != expected {(seq_len, batch_size, n_channels)}" + + loss = output.sum() + loss.backward() + + assert hidden_states.grad is not None, "Gradients should flow through offloaded path" + assert hidden_states.grad.shape == hidden_states.shape + assert hidden_states.grad.abs().sum() > 0, "Gradients should be non-trivial" + + PipelineOffloadManager.reset_instance() + + def test_offloading_numerical_equivalence(self): + """Offloaded forward+backward must produce the same result as non-offloaded. + + Compares outputs and gradients between a layer with offloading disabled + vs enabled to ensure the offloading path does not corrupt activations. + """ + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + PipelineOffloadManager, + ) + + PipelineOffloadManager.reset_instance() + + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + n_channels = num_streams * hidden_size + + torch.manual_seed(42) + input_data = torch.randn(seq_len, batch_size, n_channels, device='cuda') + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + # Run without offloading + config_no_offload = _make_mhc_config(hidden_size=hidden_size, num_streams=num_streams) + layer_spec = _make_mhc_layer_spec() + layer_no_offload = HyperConnectionTransformerLayer( + config_no_offload, layer_spec.submodules + ).cuda() + layer_no_offload.train() + + h1 = input_data.clone().detach().requires_grad_(True) + out1, _ = layer_no_offload(hidden_states=h1, attention_mask=attention_mask) + out1.sum().backward() + grad_no_offload = h1.grad.clone() + out1_detached = out1.detach().clone() + + # Run with offloading using the same weights + config_offload = _make_mhc_config( + hidden_size=hidden_size, + num_streams=num_streams, + fine_grained_activation_offloading=True, + offload_modules=["attn_norm", "mlp_norm"], + ) + layer_offload = HyperConnectionTransformerLayer( + config_offload, layer_spec.submodules + ).cuda() + layer_offload.load_state_dict(layer_no_offload.state_dict()) + layer_offload.train() + + mgr = PipelineOffloadManager.get_instance() + mgr.init_model_chunk_offload_handler( + pp_rank=0, vp_size=1, vp_stage=0, min_offloaded_tensor_size=0 + ) + + h2 = input_data.clone().detach().requires_grad_(True) + out2, _ = layer_offload(hidden_states=h2, attention_mask=attention_mask) + out2.sum().backward() + grad_offload = h2.grad.clone() + + PipelineOffloadManager.reset_instance() + + assert torch.allclose(out1_detached, out2.detach(), atol=1e-5), ( + f"Forward outputs differ: max diff = " + f"{(out1_detached - out2.detach()).abs().max().item()}" + ) + assert torch.allclose(grad_no_offload, grad_offload, atol=1e-5), ( + f"Gradients differ: max diff = " + f"{(grad_no_offload - grad_offload).abs().max().item()}" + ) diff --git a/train_rl.py b/train_rl.py index acf54680f4a..2f7000fb917 100644 --- a/train_rl.py +++ b/train_rl.py @@ -18,18 +18,20 @@ from megatron.rl.rl_utils import ( calculate_grpo_loss, get_logprobs, + get_rl_packed_seq_params_for_cuda_graph, get_rl_runtime_state, load_packed_data_by_index, ) from megatron.training import get_args, get_timers, pretrain, print_rank_0 -from megatron.training.utils import is_hybrid_model +from megatron.training.argument_utils import ( + gpt_config_from_args, + hybrid_config_from_args, + pretrain_cfg_container_from_args, +) from megatron.training.arguments import core_transformer_config_from_args, parse_and_validate_args -from megatron.training.argument_utils import gpt_config_from_args, hybrid_config_from_args, pretrain_cfg_container_from_args +from megatron.training.utils import is_hybrid_model from model_provider import model_provider -from megatron.core.packed_seq_params import PackedSeqParams -from megatron.rl.sequence_packing_utils import get_default_packed_seq_params - stimer = StragglerDetector() import logging @@ -260,22 +262,12 @@ def forward_step(data_iterator, model: GPTModel, loss_only: bool = False): model_to_use = model[0] if isinstance(model, list) else model if packed_seq_params is None: - if args.rl_use_sequence_packing: - packed_seq_params = get_default_packed_seq_params( - seq_length=tokens.shape[1], - max_sequences_per_bin=args.rl_sequence_packing_max_sequences_per_bin, - device=tokens.device, - ) - else: - cu_seqlens = torch.tensor([0, tokens.shape[1]], dtype=torch.int32, device=tokens.device) - packed_seq_params = PackedSeqParams( - qkv_format='thd', - cu_seqlens_q=cu_seqlens, - cu_seqlens_kv=cu_seqlens, - max_seqlen_q=tokens.shape[1], - max_seqlen_kv=tokens.shape[1], - total_tokens=tokens.shape[1], - ) + packed_seq_params = get_rl_packed_seq_params_for_cuda_graph( + seq_length=tokens.shape[1], + device=tokens.device, + sequence_packing=args.rl_use_sequence_packing, + max_sequences_per_bin=args.rl_sequence_packing_max_sequences_per_bin, + ) # Clear RoPE cache to avoid inference tensor errors try: