diff --git a/examples/dsv4_hybrid/README.md b/examples/dsv4_hybrid/README.md new file mode 100644 index 00000000000..f2a352f67cc --- /dev/null +++ b/examples/dsv4_hybrid/README.md @@ -0,0 +1,94 @@ +# DeepSeek-V4 hybrid attention (CSA / HCA) in the hybrid model + +This directory holds an example slurm script for pretraining a Mamba-based hybrid +model that mixes DeepSeek-V4's two new attention variants: + +* **CSA — Compressed Sparse Attention** (Section 2.3.1 of the DSv4 tech report). + Compresses every 4 tokens (with overlap = 2) into one KV entry and uses a + learned lightning indexer to pick the top-k relevant compressed entries. Each + query then attends to a small sliding window plus the selected compressed + positions. +* **HCA — Heavily-Compressed Attention** (Section 2.3.2). Compresses every 128 + tokens (no overlap) into one KV entry and applies dense attention over all + valid compressed positions, again concatenated with a sliding window. + +Both share the same module (`DSv4HybridSelfAttention` / `CompressedSparseAttention`); +the per-layer compression ratio selects the behaviour. + +## Pattern symbols + +The hybrid layer pattern grows two new symbols in addition to the existing +`M / G / * / D / - / E`: + +| Symbol | Meaning | +| ------ | --------------------------------------------------------- | +| `M` | Mamba | +| `G` | Gated DeltaNet | +| `*` | Standard self-attention (GQA) | +| `D` | DeepSeek Sparse Attention (DSA, MLA-style + indexer) | +| `C` | DSv4 Compressed Sparse Attention (CSA, ratio = 4) | +| `H` | DSv4 Heavily-Compressed Attention (HCA, ratio = 128) | +| `-` | Dense MLP | +| `E` | MoE | + +Only one of `*` and the MLA-like family (`D`, `C`, `H`) may appear in the same +model — they share the MLA-style q/kv projection setup. + +## Required CLI flags + +A model that uses `C` or `H` layers should set: + +* `--hybrid-layer-pattern` containing `C` and/or `H`. +* The MLA-related flags: `--q-lora-rank`, `--qk-head-dim`, `--qk-pos-emb-head-dim`, + `--v-head-dim`, plus `--rope-type rope|yarn`. +* DSA indexer flags (used only by `C`): `--dsa-indexer-n-heads`, + `--dsa-indexer-head-dim`, `--dsa-indexer-topk`, + optionally `--dsa-indexer-loss-coeff` for KL training of the indexer. + +CSA/HCA-specific flags: + +* `--csa-window-size N` — sliding-window length (default 128). +* `--csa-compress-ratio-for-c N` — ratio used by every `C` layer (default 4). +* `--csa-compress-ratio-for-h N` — ratio used by every `H` layer (default 128). +* `--csa-compress-ratios "[...]"` — explicit per-layer ratios. Overrides the + pattern-derived defaults; length must equal `num_layers`. +* `--csa-compress-rotary-base FLOAT` — RoPE base for compressed KV positions + (default 40000). +* `--csa-dense-mode` — run all `C` layers in dense (no-indexer) mode. Useful as + a warmup phase before sparse training. +* `--csa-no-attention-sink` — disable the per-head learnable sink logit. +* `--o-groups N`, `--o-lora-rank N` — grouped output projection geometry. + ``num_attention_heads * v_head_dim`` must be divisible by `o_groups`. + +## Prototype scope + +This is a CP=1, TP=1 prototype that uses the unfused RoPE path. MTP, packed +sequences, FP8/FP4, and fine-grained activation offloading are not supported +yet. Inference is also disabled for `C`/`H` layers. + +## Running + +``` +bash examples/dsv4_hybrid/train_dsv4_hybrid.sh +``` + +The provided script trains a 2B-class hybrid model with the pattern + +``` +M-M-MCM-MHM-MDM-MCM-MHM-MDM- +``` + +(24 layers, mostly Mamba + MLP, with CSA / HCA / DSA layers sprinkled in). + +Adjust `IMAGE`, `BASE_DIR`, `BLEND_PATH`, `TOKENIZER_MODEL_PATH`, etc. to your +environment. Tests must run inside the docker image — this is a slurm login +node so GPUs are not directly available. + +## Tests + +* `tests/unit_tests/transformer/experimental_attention_variant/test_csa_hca_hybrid.py` + — unit tests covering the helpers, the `Compressor`, `CSAIndexer` and the + `CompressedSparseAttention` core attention (CSA, HCA and window-only paths, + forward + backward). +* `tests/unit_tests/ssm/test_hybrid_layer_allocation.py` + — extended to exercise the new `C` and `H` symbols in the hybrid pattern. diff --git a/examples/dsv4_hybrid/train_dsv4_hybrid.sh b/examples/dsv4_hybrid/train_dsv4_hybrid.sh new file mode 100755 index 00000000000..7be55079d65 --- /dev/null +++ b/examples/dsv4_hybrid/train_dsv4_hybrid.sh @@ -0,0 +1,174 @@ +#!/bin/bash + +# DeepSeek-V4 hybrid model: Mamba + CSA / HCA / DSA + dense MLP pretraining example. +# +# Hybrid pattern symbols: +# M - Mamba +# C - DSv4 Compressed-Sparse Attention (CSA, ratio = 4 by default) +# H - DSv4 Heavily-Compressed Attention (HCA, ratio = 128 by default) +# D - DeepSeek Sparse Attention (DSA, MLA + indexer) +# - - dense MLP +# E - MoE +# +# Usage: +# ./dsv4_hybrid-1n.sh \ +# + +export CUDA_DEVICE_MAX_CONNECTIONS=1 +export NVTE_FWD_LAYERNORM_SM_MARGIN=16 +export NVTE_BWD_LAYERNORM_SM_MARGIN=16 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + +GPUS_PER_NODE=8 +MASTER_ADDR=localhost +MASTER_PORT=6000 +NUM_NODES=1 +NODE_RANK=0 +WORLD_SIZE=$(($GPUS_PER_NODE*$NUM_NODES)) + +CHECKPOINT_PATH=$1 # +TENSORBOARD_LOGS_PATH=$2 # +DATACACHE_PATH=$3 # +TOKENIZER_MODEL=$4 # +DATA_BLEND_PATH=$5 # + +SEQ_LEN=8192 +TRAIN_SAMPLES=36621094 +LR_WARMUP_SAMPLES=1024000 +LR_DECAY_SAMPLES=36621094 +LR_WSD_DECAY_SAMPLES=5493165 + +# 16-layer pattern: mamba bulk + interleaved CSA / HCA / DSA + dense MLP. +HYBRID_PATTERN="M-MCM-MHMDM-MHM-" + +DISTRIBUTED_ARGS=( + --nproc_per_node $GPUS_PER_NODE + --nnodes $NUM_NODES + --master_addr $MASTER_ADDR + --master_port $MASTER_PORT +) + +HYBRID_MODEL_ARGS=( + --hybrid-layer-pattern $HYBRID_PATTERN + --hidden-size 2048 + --num-attention-heads 16 + --num-query-groups 8 + --ffn-hidden-size 8192 + --kv-channels 128 + --mamba-num-heads 64 + --mamba-head-dim 64 + --mamba-state-dim 128 + --mamba-num-groups 8 + --seq-length $SEQ_LEN + --max-position-embeddings $SEQ_LEN + --position-embedding-type none + --normalization RMSNorm + --untie-embeddings-and-output-weights + --disable-bias-linear + --squared-relu + --init-method-std 0.0198 +) + +# CSA / HCA / DSA arguments (apply to C, H, and D layers in the pattern). +DSV4_HYBRID_ATTN_ARGS=( + --rope-type rope + --q-lora-rank 512 + --qk-head-dim 128 + --qk-pos-emb-head-dim 64 + --v-head-dim 128 + --csa-window-size 128 + --csa-compress-ratio-for-c 4 + --csa-compress-ratio-for-h 128 + --o-groups 4 + --o-lora-rank 128 + --dsa-indexer-n-heads 64 + --dsa-indexer-head-dim 128 + --dsa-indexer-topk 64 + --dsa-indexer-loss-coeff 0.0 +) + +TRAINING_ARGS=( + --micro-batch-size 1 + --global-batch-size 8 + --train-samples $TRAIN_SAMPLES + --lr 1.4e-3 + --min-lr 1.4e-5 + --lr-decay-style WSD + --lr-warmup-samples $LR_WARMUP_SAMPLES + --lr-decay-samples $LR_DECAY_SAMPLES + --lr-wsd-decay-style minus_sqrt + --lr-wsd-decay-samples $LR_WSD_DECAY_SAMPLES + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.95 + --clip-grad 1.0 + --attention-dropout 0.0 + --hidden-dropout 0.0 + --bf16 + --override-opt_param-scheduler +) + +MODEL_PARALLEL_ARGS=( + --tensor-model-parallel-size 1 + --pipeline-model-parallel-size 1 + --use-distributed-optimizer + --overlap-grad-reduce + --overlap-param-gather +) + +DATA_ARGS=( + --per-split-data-args-path $DATA_BLEND_PATH + --data-cache-path $DATACACHE_PATH + --tokenizer-type TikTokenizer + --tokenizer-model $TOKENIZER_MODEL + --tiktoken-pattern v2 + --num-workers 1 + --num-dataset-builder-threads 4 + --no-create-attention-mask-in-dataloader + --no-mmap-bin-files +) + +CHECKPOINT_ARGS=( + --save $CHECKPOINT_PATH + --load $CHECKPOINT_PATH + --ckpt-format torch_dist + --ckpt-fully-parallel-save + --ckpt-fully-parallel-load + --ckpt-assume-constant-structure + --no-load-rng + --save-interval 12500 + --save-retain-interval 100000 +) + +EVAL_AND_LOGGING_ARGS=( + --log-interval 10 + --log-throughput + --log-progress + --log-params-norm + --log-num-zeros-in-grad + --log-memory-interval 1000 + --logging-level 20 + --eval-interval 1000 + --eval-iters 14 + --tensorboard-dir $TENSORBOARD_LOGS_PATH +) + +MISC_ARGS=( + --seed 1234 + --rerun-mode disabled + --attention-backend flash + --disable-gloo-process-groups + --use-mcore-models + --spec megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec + --distributed-timeout-minutes 30 +) + +torchrun ${DISTRIBUTED_ARGS[@]} pretrain_hybrid.py \ + ${HYBRID_MODEL_ARGS[@]} \ + ${DSV4_HYBRID_ATTN_ARGS[@]} \ + ${TRAINING_ARGS[@]} \ + ${MODEL_PARALLEL_ARGS[@]} \ + ${DATA_ARGS[@]} \ + ${CHECKPOINT_ARGS[@]} \ + ${EVAL_AND_LOGGING_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 5b264b36302..1e812e804bc 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -336,18 +336,41 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC # For hybrid models, the layer map converts the global layer index to the # corresponding attention layer index or Mamba layer index depending on the # layer type. - attention_layer_map, dsa_layer_map, gdn_layer_map, mamba_layer_map = ( - operator.itemgetter( - Symbols.ATTENTION, Symbols.DS_ATTENTION, Symbols.GDN, Symbols.MAMBA - )(get_layer_maps_from_layer_type_list(mamba_inference_state_config.layer_type_list)) + ( + attention_layer_map, + dsa_layer_map, + csa_layer_map, + hca_layer_map, + gdn_layer_map, + mamba_layer_map, + ) = operator.itemgetter( + Symbols.ATTENTION, + Symbols.DS_ATTENTION, + Symbols.CSA_ATTENTION, + Symbols.HCA_ATTENTION, + Symbols.GDN, + Symbols.MAMBA, + )( + get_layer_maps_from_layer_type_list(mamba_inference_state_config.layer_type_list) ) if len(gdn_layer_map) > 0: raise NotImplementedError("GDN layers are not supported for inference.") - self.num_attention_layers = len(attention_layer_map) + len(dsa_layer_map) + self.num_attention_layers = ( + len(attention_layer_map) + + len(dsa_layer_map) + + len(csa_layer_map) + + len(hca_layer_map) + ) self.num_mamba_layers = len(mamba_layer_map) - self.layer_map = attention_layer_map | dsa_layer_map | mamba_layer_map + self.layer_map = ( + attention_layer_map + | dsa_layer_map + | csa_layer_map + | hca_layer_map + | mamba_layer_map + ) else: # The layer map is the identity function for pure Transformer models. self.num_attention_layers = model_config.num_layers // pp_size diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py index 5494d531e52..4b8024f83b4 100644 --- a/megatron/core/models/hybrid/hybrid_block.py +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -42,6 +42,8 @@ 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 mlp_layer: Union[ModuleSpec, type] = IdentityOp moe_layer: Union[ModuleSpec, type] = IdentityOp mtp_block_spec: Optional[ModuleSpec] = None @@ -146,6 +148,24 @@ def __init__( add_layer_offset=False, pp_layer_offset=pp_layer_offset, ) + elif layer_type == LayerSymbols.CSA_ATTENTION: + layer = build_module( + submodules.csa_layer, + config=self.config, + layer_number=layer_number, + pg_collection=pg_collection, + add_layer_offset=False, + pp_layer_offset=pp_layer_offset, + ) + elif layer_type == LayerSymbols.HCA_ATTENTION: + layer = build_module( + submodules.hca_layer, + config=self.config, + layer_number=layer_number, + pg_collection=pg_collection, + add_layer_offset=False, + pp_layer_offset=pp_layer_offset, + ) elif layer_type == LayerSymbols.MLP: layer = build_module( submodules.mlp_layer, diff --git a/megatron/core/models/hybrid/hybrid_layer_allocation.py b/megatron/core/models/hybrid/hybrid_layer_allocation.py index f1ba94ef7fa..9f90c37d3ac 100644 --- a/megatron/core/models/hybrid/hybrid_layer_allocation.py +++ b/megatron/core/models/hybrid/hybrid_layer_allocation.py @@ -18,11 +18,13 @@ class Symbols: GDN = 'G' ATTENTION = "*" DS_ATTENTION = "D" + CSA_ATTENTION = "C" + HCA_ATTENTION = "H" MLP = "-" MOE = 'E' PIPE = '|' MTP_SEPARATOR = "/" - VALID_LAYERS = {MAMBA, GDN, ATTENTION, DS_ATTENTION, MLP, MOE} + VALID_LAYERS = {MAMBA, GDN, ATTENTION, DS_ATTENTION, CSA_ATTENTION, HCA_ATTENTION, MLP, MOE} @classmethod def name_sorted_valid_layer_symbols(cls) -> list[str]: @@ -292,9 +294,12 @@ 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") + # Disallow Attention + MLA/DSA/CSA/HCA hybridity. + mla_like = {Symbols.DS_ATTENTION, Symbols.CSA_ATTENTION, Symbols.HCA_ATTENTION} + if Symbols.ATTENTION in pattern and any(s in pattern for s in mla_like): + raise ValueError( + "Not supported to have both standard Attention and MLA/DSA/CSA/HCA in one model" + ) def validate_segment_layers(segment: str) -> List[str]: @@ -320,9 +325,12 @@ 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") + # Disallow Attention + MLA/DSA/CSA/HCA hybridity. + mla_like = {Symbols.DS_ATTENTION, Symbols.CSA_ATTENTION, Symbols.HCA_ATTENTION} + if Symbols.ATTENTION in segment and any(s in segment for s in mla_like): + raise ValueError( + "Not supported to have both standard Attention and MLA/DSA/CSA/HCA 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 a34a45a32ba..7ec0f83fc24 100755 --- a/megatron/core/models/hybrid/hybrid_layer_specs.py +++ b/megatron/core/models/hybrid/hybrid_layer_specs.py @@ -25,12 +25,24 @@ ) from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules 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.dsa import ( DSAIndexer, DSAIndexerSubmodules, DSAttention, DSAttentionSubmodules, ) +from megatron.core.transformer.experimental_attention_variant.dsv4_hybrid_attention import ( + DSv4HybridSelfAttention, + DSv4HybridSelfAttentionSubmodules, +) from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.mlp import MLP, MLPSubmodules from megatron.core.transformer.multi_latent_attention import ( @@ -57,6 +69,49 @@ moe_grouped_gemm=True, ) + +def _make_dsv4_hybrid_attention_spec(): + """Build a self_attention spec for a DSv4 hybrid (CSA/HCA) layer. + + The same spec is used for both 'C' and 'H' symbols; the per-layer compress + ratio is selected at module-build time from ``config.csa_compress_ratios``. + """ + compressor_spec = ModuleSpec( + module=Compressor, + submodules=CompressorSubmodules(linear_wkv=TELinear, linear_wgate=TELinear, norm=TENorm), + ) + + indexer_spec = ModuleSpec( + module=CSAIndexer, + submodules=CSAIndexerSubmodules( + linear_wq_b=TELinear, linear_weights_proj=TELinear, 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=TELinear, + linear_q_up_proj=TEColumnParallelLinear, + linear_kv_proj=TEColumnParallelLinear, + core_attention=core_attention, + linear_proj=TERowParallelLinear, + q_layernorm=TENorm, + kv_layernorm=TENorm, + ), + ) + + +_dsv4_hybrid_self_attention_spec = _make_dsv4_hybrid_attention_spec() + # Inference-optimized MoE spec moe_inference = get_inference_optimized_moe_spec() @@ -164,6 +219,22 @@ self_attn_bda=get_bias_dropout_add, ), ), + csa_layer=ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + input_layernorm=TENorm, + self_attention=_dsv4_hybrid_self_attention_spec, + self_attn_bda=get_bias_dropout_add, + ), + ), + hca_layer=ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + input_layernorm=TENorm, + self_attention=_dsv4_hybrid_self_attention_spec, + self_attn_bda=get_bias_dropout_add, + ), + ), # Started with spec from gpt_layer_specs.py # Using the TE spec because we had problems getting the non-TE spec # working diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index 88a97ec777f..063583b941f 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -185,6 +185,27 @@ def __init__( self.mtp_pattern = parsed.mtp_pattern self.mtp_num_depths = parsed.mtp_num_depths + # If any CSA/HCA layers are present in the global pattern, derive a per-layer + # csa_compress_ratios list from the symbol so each layer module can read its + # own ratio at construction time. An explicit user-provided ratios list takes + # precedence. + if parsed.main_pattern is not None and getattr(config, "csa_compress_ratios", None) is None: + from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols + + global_pattern = parsed.main_pattern.replace(Symbols.PIPE, '') + if Symbols.CSA_ATTENTION in global_pattern or Symbols.HCA_ATTENTION in global_pattern: + ratio_for_c = getattr(config, "csa_compress_ratio_for_c", 4) + ratio_for_h = getattr(config, "csa_compress_ratio_for_h", 128) + ratios = [] + for ch in global_pattern: + if ch == Symbols.CSA_ATTENTION: + ratios.append(ratio_for_c) + elif ch == Symbols.HCA_ATTENTION: + ratios.append(ratio_for_h) + else: + ratios.append(0) + config.csa_compress_ratios = ratios + layer_type_list, layer_offset = select_pipeline_segment( parsed.main_pattern or '', self.pg_collection.pp, 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..2ff2523e709 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -0,0 +1,708 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Compressed-Sparse Attention (CSA) and Heavily-Compressed Attention (HCA). + +CSA and HCA are the two attention variants introduced in DeepSeek-V4 (Section 2.3 of +the technical report). Both share a single core-attention module +(``CompressedSparseAttention``); the per-layer ``compress_ratio`` controls behaviour: + +* ``ratio == 4`` (CSA) — overlap=2, sliding-window + learned indexer top-k over + compressed positions (sparse). +* ``ratio == 128`` (HCA) — overlap=1, sliding-window + dense attention over all + compressed positions (no indexer). +* ``ratio == 0`` — sliding-window only (no compression). + +This is a CP=1 / TP=1 prototype using the unfused RoPE path. +""" + +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.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.transformer.enums import AttnMaskType +from megatron.core.transformer.experimental_attention_variant.dsa import ( + DSAIndexerLossAutoScaler, + DSAIndexerLossLoggingHelper, + FusedDSAIndexerLoss, + fused_qk_topk_naive, + rotate_activation, +) +from megatron.core.transformer.module import MegatronModule +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: + 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: + 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) + + +# --------------------------------------------------------------------------- +# Helper functions for RoPE +# --------------------------------------------------------------------------- + + +def _apply_partial_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, +) -> torch.Tensor: + """Apply RoPE to the last ``pos_dim`` dims, leave the first ``nope_dim`` unchanged. + + Accepts both 3-D ``[seq, batch, head_dim]`` and 4-D ``[seq, batch, heads, head_dim]`` + inputs (a temporary head dim is inserted for the 3-D case). + """ + if ratio == 1: + total_seq_len = rotary_seq_len + else: + total_seq_len = rotary_seq_len * ratio + if config.rope_type == "rope": + rotary_pos_emb = rotary_pos_emb_module(total_seq_len, packed_seq=False) + mscale = 1.0 + else: + rotary_pos_emb, mscale = rotary_pos_emb_module(total_seq_len, packed_seq=False) + + if ratio > 1: + rotary_pos_emb = rotary_pos_emb[:total_seq_len:ratio][:rotary_seq_len] + + squeeze_head = 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=None, mscale=mscale, cp_group=cp_group + ) + out = torch.cat([x_nope, x_pe], dim=-1) + if squeeze_head: + out = out.squeeze(-2) + return out + + +# --------------------------------------------------------------------------- +# Sparse attention kernel (unfused, differentiable) with attention sink +# --------------------------------------------------------------------------- + + +def unfused_compressed_sparse_attn( + query: torch.Tensor, + kv_full: torch.Tensor, + attn_sink: torch.Tensor, + topk_indices: torch.Tensor, + softmax_scale: float, + use_attn_sink: bool = True, +) -> torch.Tensor: + """Differentiable sparse MQA with optional learnable per-head attention sink. + + Args: + query: ``[sq, b, np, hn]`` multi-head query. + kv_full: ``[n_kv, b, hn]`` single-head KV (original + compressed). + attn_sink: ``[np]`` per-head learnable sink logit. + topk_indices: ``[b, sq, topk]`` indices into ``kv_full`` (-1 = invalid). + softmax_scale: float + use_attn_sink: when False the sink is omitted. + + Returns: + ``[sq, b, np * hn]`` + """ + sq, b, np_, hn = query.size() + + # --- Gather KV at topk positions --- + kv_t = kv_full.permute(1, 0, 2) # [b, n_kv, hn] + + safe_indices = topk_indices.clamp(min=0).long() + safe_indices_exp = safe_indices.unsqueeze(-1).expand(-1, -1, -1, hn) + kv_gathered = torch.gather( + kv_t.unsqueeze(1).expand(-1, sq, -1, -1), dim=2, index=safe_indices_exp + ) + + # --- Attention scores --- + q = query.permute(1, 2, 0, 3).float() # [b, np, sq, hn] + kv_g = kv_gathered.float() # [b, sq, topk, hn] + + scores = torch.einsum("bnsh,bskh->bnsk", q, kv_g) * softmax_scale + + invalid_mask = (topk_indices < 0).unsqueeze(1) # [b, 1, sq, topk] + scores = scores.masked_fill(invalid_mask, float("-inf")) + + if use_attn_sink: + sink = attn_sink.view(1, np_, 1, 1).float() + scores_max = scores.max(dim=-1, keepdim=True).values + scores_max = torch.max(scores_max, sink) + # If a row has no valid KV, scores_max is still finite (sink), so safe. + exp_scores = torch.exp(scores - scores_max) + exp_sink = torch.exp(sink - scores_max) + sum_exp = exp_scores.sum(dim=-1, keepdim=True) + exp_sink + attn_weights = exp_scores / sum_exp + else: + # Without sink, fall back to standard masked softmax. + # If a row has no valid kv, set the attn weights to 0 (skip token). + all_invalid = invalid_mask.all(dim=-1, keepdim=True) # [b, 1, sq, 1] + scores = scores.masked_fill(all_invalid, 0.0) + attn_weights = torch.softmax(scores, dim=-1) + attn_weights = attn_weights.masked_fill(all_invalid, 0.0) + + output = torch.einsum("bnsk,bskh->bnsh", attn_weights, kv_g) + output = output.to(query.dtype) + output = output.permute(2, 0, 1, 3).contiguous().reshape(sq, b, np_ * hn) + return output + + +# --------------------------------------------------------------------------- +# Compressor +# --------------------------------------------------------------------------- + + +@dataclass +class CompressorSubmodules: + """Submodule specs for ``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. + + Compresses a sequence of tokens into a shorter sequence by pooling groups of + ``compress_ratio`` tokens with learned gated weights. + + For ``compress_ratio == 4`` (CSA) overlap=2 is used; for larger ratios (e.g. 128 + for HCA) overlap=1 is used. + """ + + 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, + ) -> None: + super().__init__(config=config) + + if pg_collection is None: + 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 + + 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", + ) + + 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", + ) + + # Learned positional bias B (kept fp32 for stability) + _ape = torch.empty( + compress_ratio, proj_out_dim, device=torch.cuda.current_device(), dtype=torch.float32 + ) + config.init_method(_ape) + self.ape = 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. + + ``[n_groups, ratio, b, coff * head_dim] -> [n_groups, 2 * ratio, b, head_dim]``. + """ + 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 forward(self, x: torch.Tensor) -> Optional[torch.Tensor]: + """Compress hidden states. + + Returns ``[sq // ratio, b, head_dim]`` or ``None`` if the input is shorter + than ``compress_ratio``. + """ + nvtx_range_push("compressor") + + sq, b, _ = x.size() + ratio = self.compress_ratio + + if sq < ratio: + nvtx_range_pop("compressor") + return None + + kv, _ = self.linear_wkv(x) + score, _ = self.linear_wgate(x) + + cutoff = (sq // ratio) * ratio + if cutoff < sq: + kv = kv[:cutoff] + score = score[:cutoff] + + n_compressed = cutoff // ratio + + kv = kv.view(n_compressed, ratio, b, -1) + score = score.view(n_compressed, ratio, b, -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)) + + kv = _apply_partial_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) + + nvtx_range_pop("compressor") + return kv + + +# --------------------------------------------------------------------------- +# CSA Indexer (top-k retrieval over compressed positions) +# --------------------------------------------------------------------------- + + +@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 KV positions for CSA. + + Reuses the index-score logic from ``DSAIndexer`` (einsum -> ReLU -> weight -> sum + -> top-k) 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, + ) -> None: + super().__init__(config=config) + + if pg_collection is None: + 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 + + 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", + ) + + 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", + ) + + 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, + ) + + def forward_before_topk( + self, x: torch.Tensor, qr: torch.Tensor, packed_seq_params: Optional[PackedSeqParams] = None + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute Q, compressed K, and weights before top-k selection.""" + nvtx_range_push("csa_indexer_before_topk") + + sq, bsz, _ = x.size() + + q, _ = self.linear_wq_b(qr) + q = q.reshape(sq, bsz, self.index_n_heads, self.index_head_dim) + q = _apply_partial_rope( + q, + self.index_head_dim - self.qk_pos_emb_head_dim, + self.qk_pos_emb_head_dim, + self.rotary_pos_emb, + self.config, + sq, + ratio=1, + cp_group=self.pg_collection.cp, + ) + q = rotate_activation(q) + + k = self.compressor(x) # [sq//ratio, b, index_head_dim] + + weights, _ = self.linear_weights_proj(x) + weights = weights * (self.index_n_heads**-0.5) + + nvtx_range_pop("csa_indexer_before_topk") + return q, k, 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)``.""" + nvtx_range_push("csa_indexer") + assert packed_seq_params is None, "Packed sequence is not supported for CSAIndexer." + q, k, weights = self.forward_before_topk(x, qr, packed_seq_params) + nvtx_range_push("csa_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("csa_indexer_qk_topk") + nvtx_range_pop("csa_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): + """Core attention used by both CSA and HCA layers. + + Combines sliding-window attention with compressed KV attention. Behaviour + depends on ``compress_ratio``: + + * ``ratio == 0`` : window-only. + * ``ratio == 4`` : window + 4x compressed + learned indexer (CSA). + * ``ratio > 4`` : window + ``ratio``x compressed, attend to all (HCA). + """ + + 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, + ): + super().__init__(config=config) + + if pg_collection is None: + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + self.pg_collection = pg_collection + + self.layer_number = 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.use_attention_sink = getattr(config, 'csa_attention_sink', True) + + # Learnable per-head attention sink + self.attn_sink = 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, + ) + else: + self.compressor = None + + # Conditionally build Indexer (only ratio == 4 + non-dense mode) + 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, + ) + else: + self.indexer = None + + 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, + ) -> torch.Tensor: + """Forward. + + Args: + query: ``[sq, b, np, v_head_dim]`` + key: ``[sq, b, 1, v_head_dim]`` + value: unused (key == value in MQA) + attention_mask: ignored (always causal here). + x: ``[sq, b, hidden_size]`` original hidden states. + qr: ``[sq, b, q_lora_rank]`` compressed query representation. + """ + nvtx_range_push("compressed_sparse_attn") + assert packed_seq_params is None, "Packed sequence is not supported for CSA/HCA." + + sq, b, np, hn = query.size() + + kv = key.squeeze(-2) # [sq, b, v_head_dim] + + # --- Compression --- + 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 + n_compressed = 0 + else: + kv_full = kv + n_compressed = 0 + + offset = sq # compressed positions are appended after original positions + + # --- Sliding window indices --- + window_idxs = get_window_topk_idxs(self.window_size, b, sq, query.device) + + indexer_loss = None + + if self.compress_ratio > 1 and n_compressed > 0: + nvtx_range_push("compressed_indices") + if self.indexer is not None: + # CSA: learned top-k indexer + 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) + ) + + 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) or 0.0 + key_for_loss = compressed_kv.unsqueeze(2).expand(-1, -1, np, -1) + weights_for_unfused = weights_indexer * 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, + ) + if indexer_loss_coeff > 0: + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=indexer_loss, + layer_number=self.layer_number, + num_layers=self.config.num_layers, + ) + 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 < n_valid_per_pos + compress_topk_idxs = torch.where( + valid, topk_indices_compressed + offset, torch.tensor(-1, device=x.device) + ) + else: + # HCA / CSA-dense: attend to all valid compressed positions + 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, + use_attn_sink=self.use_attention_sink, + ) + nvtx_range_pop("sparse_attn_kernel") + + if indexer_loss is not None and self.training and torch.is_grad_enabled(): + output = DSAIndexerLossAutoScaler.apply(output, indexer_loss) + + nvtx_range_pop("compressed_sparse_attn") + return output diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index 3734db7043f..1a81457e878 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -167,6 +167,7 @@ def compute_dsa_indexer_loss( loss_coeff: float, sparse_loss: bool, pg_collection: ProcessGroupCollection, + causal_mask_override: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Compute KL divergence loss between index_scores and true attention_scores. @@ -187,6 +188,10 @@ def compute_dsa_indexer_loss( sparse_loss: bool, whether to use sparse indexer loss. If True, only the topk indices will be used to compute the loss. pg_collection: Process group collection, must have TP process group. + causal_mask_override: Optional [b, sq, sk] additive mask (with -inf for invalid + positions and 0 for valid). When provided, used instead of the standard + upper-triangular causal mask. Used by CSA where compressed KV positions do + not follow the dense causal pattern. Returns: index_loss: KL divergence loss (scalar). @@ -203,29 +208,50 @@ def compute_dsa_indexer_loss( # Reshape to [b, np, sq, sk] attention_scores = attention_scores.reshape(b, np, sq, sk) - # causal_mask [sq, sk] - causal_mask = torch.triu( - torch.full((sq, sk), float('-inf'), dtype=torch.float32, device=attention_scores.device), - diagonal=1, - ) + if causal_mask_override is not None: + causal_mask = causal_mask_override.to(dtype=torch.float32) # [b, sq, sk] + else: + causal_mask = torch.triu( + torch.full( + (sq, sk), float('-inf'), dtype=torch.float32, device=attention_scores.device + ), + diagonal=1, + ) # index_mask [b, sq, sk] index_mask = torch.full( (b, sq, sk), float("-inf"), dtype=torch.float32, device=causal_mask.device ).scatter_(-1, topk_indices, 0) - # [b, np, sq, skv] + [1, 1, sq, skv] -> [b, np, sq, skv] - attention_scores += causal_mask.view(1, 1, sq, sk) + if causal_mask.dim() == 3: + attention_scores = attention_scores + causal_mask.unsqueeze(1) + else: + attention_scores = attention_scores + causal_mask.view(1, 1, sq, sk) if sparse_loss: # [b, np, sq, sk] + [b, 1, sq, sk] -> [b, np, sq, sk] attention_scores += index_mask.view(b, 1, sq, sk) # [b, sq, sk] + [b, sq, sk] -> [b, sq, sk] index_scores += index_mask + # Rows where every KV position is masked produce NaN from softmax(all -inf). Detect + # and zero out their logits before softmax, then mask their contributions out. + row_valid = (causal_mask > float('-inf')).any(dim=-1) + if row_valid.dim() == 1: + attn_row_mask = row_valid.view(1, 1, sq, 1) + idx_row_mask = row_valid.view(1, sq, 1) + else: + attn_row_mask = row_valid.view(b, 1, sq, 1) + idx_row_mask = row_valid.view(b, sq, 1) + attention_scores = attention_scores.masked_fill(~attn_row_mask, 0.0) + index_scores = index_scores.masked_fill(~idx_row_mask, 0.0) + # [b, np, sq, sk] -> [b, np, sq, sk] attention_scores = torch.nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32) # [b, sq, sk] -> [b, sq, sk] index_scores = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32) + attention_scores = attention_scores * attn_row_mask.float() + index_scores = index_scores * idx_row_mask.float() + # Sum attention scores across heads. # [batch, heads, seqlen_q, seqlen_k] -> [batch, seqlen_q, seqlen_k] attention_scores = attention_scores.sum(dim=1) @@ -234,7 +260,9 @@ def compute_dsa_indexer_loss( torch.distributed.all_reduce(attention_scores.contiguous(), group=pg_collection.tp) # L1 normalize target on the last dimension. Doesn't use abs() because attention_scores are # obtained from softmax so they are already non-negative. - attention_scores = attention_scores / attention_scores.sum(dim=-1, keepdim=True) + attention_scores = attention_scores / ( + attention_scores.sum(dim=-1, keepdim=True).clamp(min=1e-10) + ) # Compute KL divergence: KL(target || index) = target(x) * log(target(x) / index(x)) # kl_per_element [b, sq, sk] @@ -338,6 +366,7 @@ def fwd_fused_indexer_loss_naive( loss_coeff, sparse_loss, pg_collection, + causal_mask_override=mask, ) return topk_indices, indexer_loss @@ -355,6 +384,7 @@ def bwd_fused_indexer_loss_naive( sparse_loss, grad_loss, pg_collection, + causal_mask_override=None, ): """Naive implementation of backward pass for indexer loss.""" index_scores = _compute_index_scores(q, weights, k) # [B, Sq, Sk] @@ -374,23 +404,27 @@ def bwd_fused_indexer_loss_naive( # Reshape to [b, np, sq, sk] attention_scores = attention_scores.reshape(b, np, sq, sk) - # causal_mask [sq, sk] - causal_mask = torch.triu( - torch.full((sq, sk), float('-inf'), dtype=torch.float32, device=attention_scores.device), - diagonal=1, - ) + if causal_mask_override is not None: + causal_mask = causal_mask_override.to(dtype=torch.float32) # [b, sq, sk] + else: + causal_mask = torch.triu( + torch.full( + (sq, sk), float('-inf'), dtype=torch.float32, device=attention_scores.device + ), + diagonal=1, + ) # index_mask [b, sq, sk] index_mask = torch.full( (b, sq, sk), float("-inf"), dtype=torch.float32, device=causal_mask.device ).scatter_(-1, topk_indices, 0) # Apply causal mask to both attention and index scores - # [b, np, sq, skv] + [1, 1, sq, skv] -> [b, np, sq, skv] - attention_scores = attention_scores + causal_mask.view(1, 1, sq, sk) - # [b, sq, sk] + [1, sq, sk] -> [b, sq, sk] - index_scores = index_scores + causal_mask.unsqueeze(0) - # Free causal_mask - no longer needed - del causal_mask + if causal_mask.dim() == 3: + attention_scores = attention_scores + causal_mask.unsqueeze(1) + index_scores = index_scores + causal_mask + else: + attention_scores = attention_scores + causal_mask.view(1, 1, sq, sk) + index_scores = index_scores + causal_mask.unsqueeze(0) if sparse_loss: # [b, np, sq, sk] + [b, 1, sq, sk] -> [b, np, sq, sk] @@ -398,6 +432,19 @@ def bwd_fused_indexer_loss_naive( # [b, sq, sk] + [b, sq, sk] -> [b, sq, sk] index_scores = index_scores + index_mask + # Detect rows where every KV is masked (early CSA query positions); zero those rows + # before softmax to avoid NaN, then mask out the corresponding gradients afterwards. + row_valid = (causal_mask > float('-inf')).any(dim=-1) + del causal_mask + if row_valid.dim() == 1: + attn_row_mask = row_valid.view(1, 1, sq, 1) + idx_row_mask = row_valid.view(1, sq, 1) + else: + attn_row_mask = row_valid.view(b, 1, sq, 1) + idx_row_mask = row_valid.view(b, sq, 1) + attention_scores = attention_scores.masked_fill(~attn_row_mask, 0.0) + index_scores = index_scores.masked_fill(~idx_row_mask, 0.0) + # Compute softmax for both attention_scores_softmax = torch.nn.functional.softmax( attention_scores, dim=-1, dtype=torch.float32 @@ -409,6 +456,9 @@ def bwd_fused_indexer_loss_naive( # Free index_scores - no longer needed after softmax del index_scores + attention_scores_softmax = attention_scores_softmax * attn_row_mask.float() + index_scores_softmax = index_scores_softmax * idx_row_mask.float() + # Sum attention scores across heads: [b, np, sq, sk] -> [b, sq, sk] attention_scores_sum = attention_scores_softmax.sum(dim=1) # Free attention_scores_softmax @@ -421,7 +471,7 @@ def bwd_fused_indexer_loss_naive( # L1 normalize attention_scores_normalized = attention_scores_sum / attention_scores_sum.sum( dim=-1, keepdim=True - ) + ).clamp(min=1e-10) # Free attention_scores_sum - no longer needed after normalization del attention_scores_sum @@ -451,20 +501,27 @@ def bwd_fused_indexer_loss_naive( del index_scores_softmax, grad_index_scores_softmax, sum_grad # Zero out gradients for masked positions - # Create a mask for valid (non-masked) positions - # Causal mask: position (i, j) is valid if j <= i - causal_valid_mask = torch.tril( - torch.ones((sq, sk), device=q.device, dtype=torch.bool) - ) # [sq, sk] + if causal_mask_override is not None: + _cm = causal_mask_override.to(dtype=torch.float32) + if _cm.dim() == 2: + _cm = _cm.unsqueeze(0) + causal_valid_mask = (_cm == 0).squeeze(0) if _cm.shape[0] == 1 else (_cm == 0) + else: + causal_valid_mask = torch.tril(torch.ones((sq, sk), device=q.device, dtype=torch.bool)) + + if causal_valid_mask.dim() == 2: + causal_valid_mask = causal_valid_mask.unsqueeze(0) + causal_valid_mask = causal_valid_mask.expand(b, sq, sk) + if sparse_loss: # Also apply index mask - only topk positions are valid index_valid_mask = index_mask == 0 # [b, sq, sk] del index_mask # Free index_mask immediately after use - valid_mask = causal_valid_mask.unsqueeze(0) & index_valid_mask # [b, sq, sk] + valid_mask = causal_valid_mask & index_valid_mask # [b, sq, sk] del index_valid_mask else: del index_mask # Free index_mask even if not used for sparse_loss - valid_mask = causal_valid_mask.unsqueeze(0).expand(b, sq, sk) # [b, sq, sk] + valid_mask = causal_valid_mask # [b, sq, sk] del causal_valid_mask grad_index_scores_logits = grad_index_scores_logits * valid_mask.float() @@ -543,7 +600,7 @@ def forward( ) # Save for backward (recomputation strategy) - ctx.save_for_backward(q, weights, k, query, key, topk_indices) + ctx.save_for_backward(q, weights, k, query, key, topk_indices, mask) ctx.softmax_scale = softmax_scale ctx.loss_coeff = loss_coeff ctx.sparse_loss = sparse_loss @@ -556,7 +613,7 @@ def backward(ctx, grad_topk_indices, grad_loss): """ Backward: Recompute what we need. """ - q, weights, k, query, key, topk_indices = ctx.saved_tensors + q, weights, k, query, key, topk_indices, mask = ctx.saved_tensors grad_q, grad_weights, grad_k = bwd_fused_indexer_loss_naive( q, @@ -570,6 +627,7 @@ def backward(ctx, grad_topk_indices, grad_loss): ctx.sparse_loss, grad_loss, ctx.pg_collection, + causal_mask_override=mask, ) # query and key are detached in forward, so return None for their gradients diff --git a/megatron/core/transformer/experimental_attention_variant/dsv4_hybrid_attention.py b/megatron/core/transformer/experimental_attention_variant/dsv4_hybrid_attention.py new file mode 100644 index 00000000000..5d5d01e12eb --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/dsv4_hybrid_attention.py @@ -0,0 +1,366 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""DeepSeek-V4 hybrid attention layer. + +A self-attention module that wraps the ``CompressedSparseAttention`` core attention +together with MLA-style query/key projections, partial RoPE, and a grouped output +projection. + +The same module is used for CSA and HCA layers; the per-layer ``compress_ratio`` +(controlled by the layer-allocation pattern + config) selects between behaviours. + +This is a CP=1 / TP=1 prototype using the unfused RoPE path. Inference, MTP, fp8/fp4, +and fine-grained activation offloading are not supported. +""" + +from dataclasses import dataclass +from typing import Optional, Union + +import torch +import torch.nn as nn + +from megatron.core.models.common.embeddings import ( + RotaryEmbedding, + YarnRotaryEmbedding, + apply_rotary_pos_emb, +) +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_config import MLATransformerConfig + + +@torch.compile +def _q_rms_norm(q: torch.Tensor, eps: float) -> torch.Tensor: + """Fused RMS normalization for query (no learnable weight).""" + return q * torch.rsqrt(q.square().mean(-1, keepdim=True) + eps) + + +@dataclass +class DSv4HybridSelfAttentionSubmodules: + """Submodules for the ``DSv4HybridSelfAttention`` layer.""" + + q_layernorm: Union[ModuleSpec, type] = None + kv_layernorm: Union[ModuleSpec, type] = None + 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 DSv4HybridSelfAttention(MegatronModule): + """DeepSeek-V4 hybrid (CSA / HCA) self-attention layer.""" + + def __init__( + self, + config: MLATransformerConfig, + submodules: DSv4HybridSelfAttentionSubmodules, + layer_number: int, + attn_mask_type: AttnMaskType = AttnMaskType.causal, + attention_type: str = "self", + cp_comm_type: Optional[str] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + compress_ratio: Optional[int] = None, + **kwargs, + ) -> None: + super().__init__(config=config) + self.config: MLATransformerConfig + self.layer_number = layer_number + self.attn_mask_type = attn_mask_type + self.attention_type = attention_type + + if pg_collection is None: + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + self.pg_collection = pg_collection + assert self.pg_collection.tp.size() == 1, "DSv4 hybrid attention requires TP size 1." + assert self.pg_collection.cp.size() == 1, "DSv4 hybrid attention requires CP size 1." + + self.num_attention_heads_per_partition = self.config.num_attention_heads + self.q_head_dim = self.config.v_head_dim + self.query_projection_size = self.q_head_dim * self.config.num_attention_heads + + # Resolve the per-layer compression ratio. + if compress_ratio is None: + compress_ratios = self.config.csa_compress_ratios + assert compress_ratios is not None and len(compress_ratios) >= layer_number, ( + f"csa_compress_ratios must be set and have length >= num_layers; " + f"got {compress_ratios} for layer {layer_number}" + ) + compress_ratio = compress_ratios[layer_number - 1] + self.compress_ratio = compress_ratio + + # =========================================================================== + # RoPE + # =========================================================================== + rope_base = self.config.rotary_base + if compress_ratio > 1: + rope_base = self.config.csa_compress_rotary_base + if self.config.rope_type == "rope": + 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, + ) + elif self.config.rope_type == "yarn": + 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, + ) + else: + raise ValueError( + f"Unsupported RoPE type: {self.config.rope_type}, supported types are " + "'rope' and 'yarn'" + ) + + # =========================================================================== + # QKV projections + # =========================================================================== + 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, + parallel_mode='duplicated', + ) + + 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=self.pg_collection.tp, + ) + + 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_proj', + tp_group=self.pg_collection.tp, + ) + + self.q_layernorm = build_module( + submodules.q_layernorm, + hidden_size=self.config.q_lora_rank, + config=self.config, + eps=self.config.layernorm_epsilon, + ) + self.kv_layernorm = build_module( + submodules.kv_layernorm, + hidden_size=self.config.v_head_dim, + config=self.config, + eps=self.config.layernorm_epsilon, + ) + + # =========================================================================== + # Core attention (CSA / HCA) + # =========================================================================== + 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=None, + k_channels=self.q_head_dim, + v_channels=self.config.v_head_dim, + cp_comm_type=cp_comm_type, + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, + ) + + # =========================================================================== + # Grouped output projection + # =========================================================================== + 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 = 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, + ) + + # =========================================================================== + # Helpers + # =========================================================================== + + def _build_rotary_pos_emb(self, seqlen: int, dtype: torch.dtype): + if self.config.rope_type == "rope": + rotary_pos_emb = self.rotary_pos_emb(seqlen, packed_seq=False) + mscale = 1.0 + else: + rotary_pos_emb, mscale = self.rotary_pos_emb(seqlen, packed_seq=False) + return rotary_pos_emb, mscale + + def _get_qkv(self, hidden_states: torch.Tensor): + """Compute query, key, value, and the compressed q representation.""" + sq, b, _ = hidden_states.size() + rotary_pos_emb, mscale = self._build_rotary_pos_emb(sq, hidden_states.dtype) + + q_compressed, _ = self.linear_q_down_proj(hidden_states) + q_compressed = self.q_layernorm(q_compressed) + + q, _ = self.linear_q_up_proj(q_compressed) + q = q.view(sq, b, self.num_attention_heads_per_partition, self.q_head_dim) + q = _q_rms_norm(q, self.config.layernorm_epsilon) + + kv, _ = self.linear_kv_proj(hidden_states) + kv = self.kv_layernorm(kv) + + # Partial RoPE on the last qk_pos_emb_head_dim dims of q and kv. + pos_dim = self.config.qk_pos_emb_head_dim + q_no_pe, q_pe = torch.split(q, [self.q_head_dim - pos_dim, pos_dim], dim=-1) + q_pe = apply_rotary_pos_emb( + q_pe, + rotary_pos_emb, + config=self.config, + cu_seqlens=None, + mscale=mscale, + cp_group=self.pg_collection.cp, + ) + query = torch.cat([q_no_pe, q_pe], dim=-1) + + kv_no_pe, k_pe = torch.split(kv, [self.config.v_head_dim - pos_dim, pos_dim], dim=-1) + k_pe = apply_rotary_pos_emb( + k_pe.unsqueeze(-2), + rotary_pos_emb, + config=self.config, + cu_seqlens=None, + mscale=mscale, + cp_group=self.pg_collection.cp, + ).squeeze(-2) + kv = torch.cat([kv_no_pe, k_pe], dim=-1).unsqueeze(-2) + + return query.contiguous(), kv.contiguous(), kv.contiguous(), q_compressed + + def _inverse_partial_rope_on_output(self, core_attn_out: torch.Tensor) -> torch.Tensor: + """Apply RoPE with negated frequencies to the last ``qk_pos_emb_head_dim`` dims. + + Required because the same compressed entries are used as both keys and values, + which causes an absolute-position bias on the output. Inverse RoPE cancels it. + """ + seq_len, b, _ = core_attn_out.size() + 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, b, n_heads, -1) + + rotary_pos_emb, mscale = self._build_rotary_pos_emb(seq_len, core_attn_out.dtype) + # Inverse rotation: negate the frequencies. + rotary_pos_emb = -rotary_pos_emb + + content_part, rot_part = torch.split(core_attn_out, [nope_dim, pos_dim], dim=-1) + rot_part = apply_rotary_pos_emb( + rot_part, + rotary_pos_emb, + config=self.config, + cu_seqlens=None, + mscale=mscale, + cp_group=self.pg_collection.cp, + ) + core_attn_out = torch.cat([content_part, rot_part], dim=-1) + return core_attn_out.view(seq_len, b, -1) + + # =========================================================================== + # Forward + # =========================================================================== + + 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. Returns ``(output, bias)`` matching standard self-attention.""" + assert rotary_pos_emb is None, "DSv4HybridSelfAttention computes RoPE internally." + assert attention_bias is None, "Attention bias is not supported." + assert packed_seq_params is None, "Packed sequence is not supported." + assert ( + inference_context is None and inference_params is None + ), "Inference is not supported for DSv4HybridSelfAttention." + + query, key, value, q_compressed = self._get_qkv(hidden_states) + + core_attn_out = self.core_attention( + query, key, value, attention_mask, x=hidden_states, qr=q_compressed + ) + + core_attn_out = self._inverse_partial_rope_on_output(core_attn_out) + + # Grouped output projection: split heads into groups, project per-group, concat. + sq, b, _ = core_attn_out.size() + core_attn_out = core_attn_out.view(sq, b, self.o_local_groups, -1) + wo_a_weight = self.linear_o_group_proj.view( + self.o_local_groups, self.config.o_lora_rank, -1 + ) + # [...g d] @ [g r d]^T -> [...g r] + 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, bias = self.linear_proj(core_attn_out) + return output, bias diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index 601ae89fae1..60dc089ad5d 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -375,9 +375,9 @@ def forward( else: if inference_context is None or inference_context.is_static_batching(): extra_kwargs = {} - if self.config.experimental_attention_variant == "dsa": - # For dsa we need to pass in the original hidden states and the compressed - # query representation. + if self.config.experimental_attention_variant in ("dsa", "dsv4_hybrid"): + # DSA layers (the only path through MLASelfAttention for either variant) + # require the original hidden states and compressed query representation. extra_kwargs["x"] = hidden_states extra_kwargs["qr"] = q_compressed with off_interface( diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 40c1a745493..dac15564389 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -262,8 +262,12 @@ 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 supports gated_delta_net, dsa, and + dsv4_hybrid. ``dsv4_hybrid`` enables DeepSeek-V4 CSA/HCA layers in the hybrid model + (selected per-layer via the hybrid layer pattern).""" #################### # DSA @@ -2336,6 +2340,41 @@ class MLATransformerConfig(TransformerConfig): Otherwise fall back to the unfused MLA. """ + #################### + # DeepSeek-v4 hybrid attention (CSA / HCA) + #################### + csa_window_size: int = 128 + """Sliding-window size used by CSA / HCA core attention.""" + + csa_compress_ratios: Optional[List[int]] = None + """Per-layer compression ratios for CSA / HCA. The list length must equal + ``num_layers``. Valid values: 0 (window only), 4 (CSA), 128 (HCA). For pure CSA / HCA + layers in a hybrid model the entry is normally derived from the layer-allocation + pattern but may also be set here directly.""" + + csa_compress_rotary_base: float = 40000.0 + """RoPE base for compressed KV positions in CSA / HCA.""" + + csa_dense_mode: bool = False + """If True, the CSA layer attends to all valid compressed positions and disables the + learned indexer (used as a warmup phase before sparse training).""" + + csa_attention_sink: bool = True + """Whether to use the per-head learnable attention sink in CSA / HCA core attention.""" + + csa_compress_ratio_for_c: int = 4 + """Default compression ratio for layers marked 'C' (CSA) in the hybrid pattern.""" + + csa_compress_ratio_for_h: int = 128 + """Default compression ratio for layers marked 'H' (HCA) in the hybrid pattern.""" + + o_groups: int = 1 + """Number of groups for the grouped output projection used by DSv4 hybrid attention.""" + + o_lora_rank: int = 128 + """Per-group lora rank for the grouped output projection used by DSv4 hybrid + attention. The intermediate output is ``o_groups * o_lora_rank`` wide.""" + def __post_init__(self): super().__post_init__() if self.multi_latent_attention and self.apply_rope_fusion and self.rope_type != "yarn": diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 5f1968dcc27..93bd8aac636 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -6,46 +6,45 @@ import dataclasses import json import os -from pathlib import Path import re import types +from pathlib import Path import torch import torch.nn.functional as F from packaging.version import Version as PkgVersion +from megatron.core.activations import squared_relu from megatron.core.dist_checkpointing.validation import StrictHandling +from megatron.core.fusions.fused_bias_geglu import quick_gelu +from megatron.core.msc_utils import MultiStorageClientFeature +from megatron.core.quantization.utils import ( + kitchen_quantization_recipe_config, + load_quantization_recipe, +) from megatron.core.rerun_state_machine import RerunStateMachine from megatron.core.transformer import MLATransformerConfig, TransformerConfig -from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout from megatron.core.transformer.enums import AttnBackend, CudaGraphScope from megatron.core.transformer.heterogeneous.heterogeneous_config import ( HeterogeneousTransformerConfig, MLPConfig, ) +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.core.activations import squared_relu -from megatron.core.fusions.fused_bias_geglu import quick_gelu +from megatron.training.argument_utils import ArgumentGroupFactory 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.core.quantization.utils import ( - kitchen_quantization_recipe_config, - load_quantization_recipe, -) -from megatron.training.argument_utils import ArgumentGroupFactory def add_megatron_arguments(parser: argparse.ArgumentParser): """"Add Megatron-LM arguments to the given parser.""" @@ -336,8 +335,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 @@ -664,8 +664,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 @@ -806,8 +808,12 @@ 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: + # Infer use of MLA from unified pattern (DSA, CSA and HCA all require MLA config) + if args.hybrid_layer_pattern and ( + Symbols.DS_ATTENTION in args.hybrid_layer_pattern + or Symbols.CSA_ATTENTION in args.hybrid_layer_pattern + or Symbols.HCA_ATTENTION in args.hybrid_layer_pattern + ): args.multi_latent_attention = True # === End of hybrid layer pattern: deprecation handling and validation === @@ -1783,6 +1789,11 @@ def core_transformer_config_from_args(args, config_class=None): 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' + if ( + Symbols.CSA_ATTENTION in args.hybrid_layer_pattern + or Symbols.HCA_ATTENTION in args.hybrid_layer_pattern + ): + kw_args['experimental_attention_variant'] = 'dsv4_hybrid' kw_args['inference_sampling_seed'] = args.seed @@ -2509,8 +2520,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") @@ -3244,6 +3254,35 @@ 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.') + # DeepSeek-V4 hybrid (CSA / HCA) + group.add_argument('--csa-window-size', type=int, default=128, + help='Sliding-window size for CSA / HCA core attention.') + group.add_argument('--csa-compress-ratios', type=la_freq_type, default=None, + help='Per-layer compression ratios for CSA / HCA. Accepts a Python ' + 'list expression, e.g. "[0,0,4,128,4,128]". The list length must ' + 'equal num_layers. Valid values: 0 (window only), 4 (CSA), ' + '128 (HCA). When omitted in a hybrid model with C/H pattern ' + 'symbols, it is auto-derived from the pattern.') + group.add_argument('--csa-compress-rotary-base', type=float, default=40000.0, + help='RoPE base for compressed KV positions in CSA / HCA.') + group.add_argument('--csa-dense-mode', action='store_true', + help='If set, the CSA layer attends to all valid compressed positions ' + 'and disables the learned indexer (warmup phase).') + group.add_argument('--csa-no-attention-sink', dest='csa_attention_sink', + action='store_false', + help='Disable the learnable per-head attention sink in CSA / HCA core ' + 'attention.') + group.set_defaults(csa_attention_sink=True) + group.add_argument('--csa-compress-ratio-for-c', type=int, default=4, + help="Compression ratio for layers marked 'C' (CSA) in the hybrid pattern.") + group.add_argument('--csa-compress-ratio-for-h', type=int, default=128, + help="Compression ratio for layers marked 'H' (HCA) in the hybrid pattern.") + group.add_argument('--o-groups', type=int, default=1, + help='Number of groups for the grouped output projection used by ' + 'DSv4 hybrid attention.') + group.add_argument('--o-lora-rank', type=int, default=128, + help='Per-group lora rank for the grouped output projection used by ' + 'DSv4 hybrid attention.') return parser def _add_heterogeneous_args(parser): @@ -3310,10 +3349,10 @@ def _add_experimental_args(parser): '`transformer_block.py`, or `transformer_layer.py`') group.add_argument('--hybrid-layer-pattern', type=str, default=None, help='Specify a hybrid layer pattern using M (mamba), G (gdn), ' - '* (attention), D (dsa), - (mlp), E (moe). Use | to define pipeline ' - 'stage boundaries for flexible virtual pipeline parallel (fVPP). ' - 'Use / to separate MTP patterns. ' - 'Example: "M-M-|M-M*-|M-M-|M-M*-" or "M-M-|M-M*-/MM/MM". ' + '* (attention), D (dsa), C (DSv4 CSA), H (DSv4 HCA), - (mlp), E (moe). ' + 'Use | to define pipeline stage boundaries for flexible virtual ' + 'pipeline parallel (fVPP). Use / to separate MTP patterns. ' + 'Example: "M-M-|M-M*-|M-M-|M-M*-" or "M-MCMHM-/MM/MM". ' 'When this flag is used, it is the sole indicator that a hybrid model ' 'is being run.') group.add_argument('--hybrid-override-pattern', type=str, default=None, @@ -3372,7 +3411,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 diff --git a/tests/unit_tests/ssm/test_hybrid_layer_allocation.py b/tests/unit_tests/ssm/test_hybrid_layer_allocation.py index fe0d7c2dc1e..7cfd07d38cd 100644 --- a/tests/unit_tests/ssm/test_hybrid_layer_allocation.py +++ b/tests/unit_tests/ssm/test_hybrid_layer_allocation.py @@ -301,78 +301,73 @@ def test_dataclass_equality(self): assert p1 == p2 +def _expected_counts(**overrides): + """Build the expected ``get_hybrid_layer_counts`` dict. + + All valid layer symbols are present in the result; the keys not given as + overrides default to 0. + """ + base = {'*': 0, 'C': 0, 'D': 0, 'G': 0, 'H': 0, 'M': 0, '-': 0, 'E': 0} + base.update(overrides) + return base + + @pytest.mark.internal 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_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_counts( + **{'*': 1, 'G': 1, 'M': 1, '-': 1, 'E': 1} + ) + assert get_hybrid_layer_counts("MGD-E") == _expected_counts( + **{'D': 1, 'G': 1, 'M': 1, '-': 1, 'E': 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_counts(**{'*': 2, 'M': 2}) + assert get_hybrid_layer_counts("M-M-|M-M*-") == _expected_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_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_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_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_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_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_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_counts(**{'D': 2, 'M': 2}) + + def test_csa_hca_pattern(self): + # CSA + HCA + DSA + Mamba + assert get_hybrid_layer_counts("MCMHMD") == _expected_counts( + **{'C': 1, 'D': 1, 'H': 1, 'M': 3} + ) 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_counts() @pytest.mark.internal @@ -645,7 +640,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) == 8 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/transformer/experimental_attention_variant/test_csa_hca_hybrid.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_hca_hybrid.py new file mode 100644 index 00000000000..f3db7ac9629 --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_hca_hybrid.py @@ -0,0 +1,437 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for the CSA/HCA hybrid attention modules.""" + +from unittest.mock import patch + +import pytest +import torch + +from megatron.core.models.hybrid.hybrid_layer_allocation import ( + Symbols, + parse_hybrid_pattern, + validate_segment_layers, +) +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.experimental_attention_variant.csa import ( + CompressedSparseAttention, + CompressedSparseAttentionSubmodules, + Compressor, + CompressorSubmodules, + CSAIndexer, + CSAIndexerSubmodules, + _get_compress_topk_idxs_cached, + _get_window_topk_idxs_cached, + get_compress_topk_idxs, + get_window_topk_idxs, + unfused_compressed_sparse_attn, +) +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_config import MLATransformerConfig +from tests.unit_tests.test_utilities import Utils + +# --------------------------------------------------------------------------- +# Hadamard (rotate_activation) shim for environments without the kernel. +# --------------------------------------------------------------------------- + +try: + from fast_hadamard_transform import hadamard_transform as _hadamard_transform # noqa: F401 + + HAVE_HADAMARD = True +except ImportError: + HAVE_HADAMARD = False + + +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(): + if not HAVE_HADAMARD: + with patch( + 'megatron.core.transformer.experimental_attention_variant.dsa.hadamard_transform', + _mock_hadamard_transform, + ): + yield + else: + yield + + +# --------------------------------------------------------------------------- +# Helper-only tests (do not require model-parallel initialization). +# --------------------------------------------------------------------------- + + +@pytest.mark.internal +class TestLayerSymbols: + def test_csa_hca_symbols_present(self): + assert Symbols.CSA_ATTENTION == "C" + assert Symbols.HCA_ATTENTION == "H" + assert Symbols.CSA_ATTENTION in Symbols.VALID_LAYERS + assert Symbols.HCA_ATTENTION in Symbols.VALID_LAYERS + + def test_pattern_with_csa_hca(self): + layers = validate_segment_layers("MCMHM-") + assert layers == ['M', 'C', 'M', 'H', 'M', '-'] + + def test_csa_hca_not_combinable_with_standard_attention(self): + with pytest.raises(ValueError): + validate_segment_layers("M*MC") + with pytest.raises(ValueError): + validate_segment_layers("M*MH") + + def test_parse_pattern_with_csa_hca(self): + parsed = parse_hybrid_pattern("MCMHMD") + assert parsed.main_pattern == "MCMHMD" + assert parsed.mtp_pattern is None + + +@pytest.mark.internal +class TestIndexHelpers: + def test_window_indices_shape_and_validity(self): + idxs = _get_window_topk_idxs_cached(window_size=4, seqlen=8, device_str="cpu") + assert idxs.shape == (8, 4) + # First row only has position 0 visible (window starts before seq) + assert idxs[0, 0].item() == 0 + assert idxs[0, 1].item() == -1 + # Last row sees positions 4..7 + assert idxs[7].tolist() == [4, 5, 6, 7] + + def test_window_indices_batch_expand(self): + idxs = get_window_topk_idxs( + window_size=2, batch_size=3, seqlen=4, device=torch.device("cpu") + ) + assert idxs.shape == (3, 4, 2) + assert torch.equal(idxs[0], idxs[1]) + + def test_compress_indices_causal(self): + # ratio=4, seq=8 -> 2 compressed positions + idxs = _get_compress_topk_idxs_cached(ratio=4, seqlen=8, offset=8, device_str="cpu") + assert idxs.shape == (8, 2) + # Compressed entry j is visible to query position i once (i+1)//ratio > j. + # Positions 0..2 cannot see entry 0 yet ((i+1)//4 == 0 for i in 0..2). + assert (idxs[0:3] == -1).all().item() + # Positions 3..7 see entry 0 (offset=8 -> index 8). + assert (idxs[3:8, 0] == 8).all().item() + # Entry 1 is only visible at i=7 (where (7+1)//4 == 2 > 1); index 8+1=9. + assert (idxs[0:7, 1] == -1).all().item() + assert idxs[7, 1].item() == 9 + + def test_compress_indices_batch_expand(self): + idxs = get_compress_topk_idxs( + ratio=4, batch_size=2, seqlen=8, offset=0, device=torch.device("cpu") + ) + assert idxs.shape == (2, 8, 2) + + +@pytest.mark.internal +class TestUnfusedAttnCpuShape: + """Sanity-check the unfused sparse attn kernel on CPU (no CUDA). + + Verifies output shape and that an all-invalid row produces zero output when + ``use_attn_sink=False``. + """ + + def test_output_shape(self): + sq, b, np_, hn = 4, 2, 3, 16 + topk = 5 + query = torch.randn(sq, b, np_, hn) + kv_full = torch.randn(7, b, hn) + attn_sink = torch.zeros(np_) + topk_idxs = torch.full((b, sq, topk), -1, dtype=torch.int32) + # Make some valid positions + topk_idxs[:, :, 0] = 0 + topk_idxs[:, 1:, 1] = 1 + + out = unfused_compressed_sparse_attn( + query, kv_full, attn_sink, topk_idxs, softmax_scale=hn**-0.5, use_attn_sink=False + ) + assert out.shape == (sq, b, np_ * hn) + + def test_all_invalid_row_no_sink_yields_zero(self): + sq, b, np_, hn = 2, 1, 1, 4 + query = torch.randn(sq, b, np_, hn) + kv_full = torch.randn(2, b, hn) + attn_sink = torch.zeros(np_) + topk_idxs = torch.full((b, sq, 3), -1, dtype=torch.int32) + # Row 0 all invalid; row 1 has one valid position + topk_idxs[:, 1, 0] = 0 + + out = unfused_compressed_sparse_attn( + query, kv_full, attn_sink, topk_idxs, softmax_scale=hn**-0.5, use_attn_sink=False + ) + # Row 0 should be exactly zero + assert torch.allclose(out[0], torch.zeros_like(out[0])) + + +# --------------------------------------------------------------------------- +# CUDA-required tests +# --------------------------------------------------------------------------- + + +def _make_mla_config(num_layers: int = 2, csa_compress_ratios=None): + return MLATransformerConfig( + num_layers=num_layers, + hidden_size=64, + num_attention_heads=4, + num_query_groups=4, + kv_channels=64, + ffn_hidden_size=128, + q_lora_rank=32, + kv_lora_rank=32, + qk_head_dim=32, + qk_pos_emb_head_dim=16, + v_head_dim=32, + rope_type="rope", + rotary_base=10000, + normalization="RMSNorm", + layernorm_epsilon=1e-5, + # CSA / HCA defaults + csa_window_size=8, + csa_compress_ratios=csa_compress_ratios, + csa_dense_mode=False, + csa_attention_sink=True, + # Indexer + dsa_indexer_n_heads=2, + dsa_indexer_head_dim=32, + dsa_indexer_topk=4, + dsa_indexer_loss_coeff=0.0, + # Output projection + o_groups=2, + o_lora_rank=16, + params_dtype=torch.bfloat16, + bf16=True, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.internal +class TestCompressorAndIndexer: + @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 _build_pg(self): + return ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + def _build_rotary(self, config): + from megatron.core.models.common.embeddings import RotaryEmbedding + + pg = self._build_pg() + return RotaryEmbedding( + config.qk_pos_emb_head_dim, + rotary_percent=1.0, + rotary_base=config.rotary_base, + cp_group=pg.cp, + ) + + def _build_compressor_spec(self): + from megatron.core.extensions.transformer_engine import TELinear, TENorm + + return ModuleSpec( + module=Compressor, + submodules=CompressorSubmodules( + linear_wkv=TELinear, linear_wgate=TELinear, norm=TENorm + ), + ) + + def test_compressor_csa_overlap_ratio_4(self): + from megatron.core.transformer.spec_utils import build_module + + config = _make_mla_config(num_layers=1) + rotary = self._build_rotary(config) + + compressor = ( + build_module( + self._build_compressor_spec(), + config=config, + compress_ratio=4, + head_dim=config.v_head_dim, + rotate=False, + rotary_pos_emb=rotary, + pg_collection=self._build_pg(), + ) + .cuda() + .to(torch.bfloat16) + ) + + sq, b = 16, 2 + x = torch.randn(sq, b, config.hidden_size, device="cuda", dtype=torch.bfloat16) + out = compressor(x) + # Ratio 4 -> sq/4 compressed entries + assert out.shape == (sq // 4, b, config.v_head_dim) + + def test_compressor_hca_no_overlap_ratio_8(self): + from megatron.core.transformer.spec_utils import build_module + + config = _make_mla_config(num_layers=1) + rotary = self._build_rotary(config) + + compressor = ( + build_module( + self._build_compressor_spec(), + config=config, + compress_ratio=8, + head_dim=config.v_head_dim, + rotate=False, + rotary_pos_emb=rotary, + pg_collection=self._build_pg(), + ) + .cuda() + .to(torch.bfloat16) + ) + + sq, b = 16, 2 + x = torch.randn(sq, b, config.hidden_size, device="cuda", dtype=torch.bfloat16) + out = compressor(x) + assert out.shape == (sq // 8, b, config.v_head_dim) + + def test_compressor_returns_none_when_seq_too_short(self): + from megatron.core.transformer.spec_utils import build_module + + config = _make_mla_config(num_layers=1) + rotary = self._build_rotary(config) + + compressor = ( + build_module( + self._build_compressor_spec(), + config=config, + compress_ratio=8, + head_dim=config.v_head_dim, + rotate=False, + rotary_pos_emb=rotary, + pg_collection=self._build_pg(), + ) + .cuda() + .to(torch.bfloat16) + ) + + x = torch.randn(4, 1, config.hidden_size, device="cuda", dtype=torch.bfloat16) + assert compressor(x) is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.internal +class TestCompressedSparseAttention: + @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 _build_attention(self, compress_ratio: int): + from megatron.core.extensions.transformer_engine import TELinear, TENorm + from megatron.core.models.common.embeddings import RotaryEmbedding + from megatron.core.transformer.enums import AttnMaskType + from megatron.core.transformer.spec_utils import build_module + + config = _make_mla_config(num_layers=1) + pg = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + rotary = RotaryEmbedding( + config.qk_pos_emb_head_dim, + rotary_percent=1.0, + rotary_base=config.rotary_base, + cp_group=pg.cp, + ) + + compressor_spec = ModuleSpec( + module=Compressor, + submodules=CompressorSubmodules( + linear_wkv=TELinear, linear_wgate=TELinear, norm=TENorm + ), + ) + indexer_spec = ModuleSpec( + module=CSAIndexer, + submodules=CSAIndexerSubmodules( + linear_wq_b=TELinear, linear_weights_proj=TELinear, compressor=compressor_spec + ), + ) + + attention = ( + build_module( + ModuleSpec( + module=CompressedSparseAttention, + submodules=CompressedSparseAttentionSubmodules( + compressor=compressor_spec, indexer=indexer_spec + ), + ), + config=config, + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type="self", + softmax_scale=None, + k_channels=config.v_head_dim, + v_channels=config.v_head_dim, + cp_comm_type="p2p", + pg_collection=pg, + rotary_pos_emb=rotary, + compress_ratio=compress_ratio, + ) + .cuda() + .to(torch.bfloat16) + ) + + return config, attention + + def _make_inputs(self, config, sq=16, b=2): + np_ = config.num_attention_heads + hn = config.v_head_dim + query = torch.randn(sq, b, np_, hn, device="cuda", dtype=torch.bfloat16) + key = torch.randn(sq, b, 1, hn, device="cuda", dtype=torch.bfloat16) + value = key.clone() + x = torch.randn(sq, b, config.hidden_size, device="cuda", dtype=torch.bfloat16) + qr = torch.randn(sq, b, config.q_lora_rank, device="cuda", dtype=torch.bfloat16) + return query, key, value, x, qr + + def test_csa_forward_shape(self): + config, attn = self._build_attention(compress_ratio=4) + query, key, value, x, qr = self._make_inputs(config) + + out = attn(query, key, value, attention_mask=None, x=x, qr=qr) + sq, b, np_, hn = query.shape + assert out.shape == (sq, b, np_ * hn) + assert attn.indexer is not None # CSA path + assert attn.compressor is not None + + def test_hca_forward_shape(self): + config, attn = self._build_attention(compress_ratio=8) + query, key, value, x, qr = self._make_inputs(config) + + out = attn(query, key, value, attention_mask=None, x=x, qr=qr) + sq, b, np_, hn = query.shape + assert out.shape == (sq, b, np_ * hn) + # HCA -> no indexer, dense over compressed positions + assert attn.indexer is None + assert attn.compressor is not None + + def test_window_only_no_compression(self): + config, attn = self._build_attention(compress_ratio=0) + query, key, value, x, qr = self._make_inputs(config) + + out = attn(query, key, value, attention_mask=None, x=x, qr=qr) + sq, b, np_, hn = query.shape + assert out.shape == (sq, b, np_ * hn) + assert attn.indexer is None + assert attn.compressor is None + + def test_csa_backward_runs(self): + config, attn = self._build_attention(compress_ratio=4) + query, key, value, x, qr = self._make_inputs(config) + # Make x and qr require grad so that the indexer-loss path is exercised. + x = x.detach().requires_grad_(True) + qr = qr.detach().requires_grad_(True) + query = query.detach().requires_grad_(True) + key = key.detach().requires_grad_(True) + + attn.train() + out = attn(query, key, value, attention_mask=None, x=x, qr=qr) + out.sum().backward() + assert query.grad is not None + assert key.grad is not None