From 3f5662ff90fd55f360d99962fdeaa0436bfecffc Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Wed, 2 Sep 2026 23:06:22 +0300 Subject: [PATCH 1/6] fix: DSA indexer FP8 precision split for wk/weights_proj MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fix DSA indexer FP8 precision split: run `linear_wk` and `linear_weights_proj` in BF16 while keeping `linear_wq_b` in FP8 under hybrid FP8 training. The DSA indexer has three projection linears: `linear_wq_b` (q-projection), `linear_wk` (k-projection), and `linear_weights_proj` (index score projection). Under FP8 hybrid training, all three were quantized to FP8 by default. However, `linear_wk` and `linear_weights_proj` feed the sparse-attention index scores directly — FP8 quantization noise in these projections perturbs the top-k selection and cascades into train/inference divergence. `linear_wq_b` (the q-projection) is a standard GEMM and can safely remain FP8. Signed-off-by: Hollow Man --- .../experimental_attention_variant/dsa.py | 102 ++++++++++-------- 1 file changed, 55 insertions(+), 47 deletions(-) diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index bfc2e885cc0..b948a0662cd 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -7,6 +7,7 @@ import torch +from megatron.core.fp8_utils import get_fp8_disabled_context from megatron.core.models.common.embeddings import ( RotaryEmbedding, YarnRotaryEmbedding, @@ -1303,6 +1304,8 @@ def __init__( f'"yarn"' ) + # Indexer precision split: linear_wq_b runs FP8 (q-projection), while + # linear_wk + linear_weights_proj run BF16 (they feed index scores directly). self.linear_wq_b = build_module( submodules.linear_wq_b, self.q_lora_rank, @@ -1315,17 +1318,19 @@ def __init__( parallel_mode="duplicated", ) - self.linear_wk = build_module( - submodules.linear_wk, - self.hidden_size, - self.index_head_dim, - config=self.config, - init_method=self.config.init_method, - bias=False, - skip_bias_add=False, - skip_weight_param_allocation=False, - parallel_mode="duplicated", - ) + # wk + weights_proj run in BF16; they feed index scores directly. + with get_fp8_disabled_context(self.config, is_init=True): + self.linear_wk = build_module( + submodules.linear_wk, + self.hidden_size, + self.index_head_dim, + config=self.config, + init_method=self.config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + ) k_norm_config = copy.copy(self.config) k_norm_config.normalization = "LayerNorm" @@ -1338,17 +1343,18 @@ def __init__( submodules.k_norm, config=k_norm_config, hidden_size=self.index_head_dim, eps=k_norm_eps ) - self.linear_weights_proj = build_module( - submodules.linear_weights_proj, - self.hidden_size, - self.index_n_heads, - config=self.config, - init_method=self.config.init_method, - bias=False, - skip_bias_add=False, - skip_weight_param_allocation=False, - parallel_mode="duplicated", - ) + with get_fp8_disabled_context(self.config, is_init=True): + self.linear_weights_proj = build_module( + submodules.linear_weights_proj, + self.hidden_size, + self.index_n_heads, + config=self.config, + init_method=self.config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + ) # Indexer projections are duplicated across tensor-parallel ranks, so their gradients # should be averaged during final gradient synchronization. for param in self.parameters(): @@ -1427,7 +1433,7 @@ def forward_before_topk( seqlen, bsz, _ = x.size() # ========================================= - # q linear and apply rope to q + # q linear and apply rope to q (FP8) # ========================================= # [seqlen, batch, q_lora_rank] -> [seqlen, batch, index_n_heads * index_head_dim] q, _ = self.linear_wq_b(qr) @@ -1435,35 +1441,37 @@ def forward_before_topk( # -> [seqlen, batch, index_n_heads, index_head_dim] q = q.reshape(seqlen, bsz, self.index_n_heads, self.index_head_dim) q = self._apply_rope(q, rotary_pos_emb, mscale, cu_seqlens=cu_seqlens_q) - - # ========================================= - # k linear and apply rope to k - # ========================================= - # [seqlen, batch, hidden_size] -> [seqlen, batch, index_head_dim] - k, _ = self.linear_wk(x) - if self.config.dsa_indexer_k_norm_fp32: - k_dtype = k.dtype - k = self.k_norm(k.float()).to(dtype=k_dtype) - else: - k = self.k_norm(k) - # [seqlen, batch, index_head_dim] -> [seqlen, batch, 1, index_head_dim] - k = k.reshape(seqlen, bsz, 1, self.index_head_dim) - k = self._apply_rope(k, rotary_pos_emb, mscale, cu_seqlens=cu_seqlens_kv) - # [seqlen, batch, 1, index_head_dim] -> [seqlen, batch, index_head_dim] - k = k.reshape(seqlen, bsz, self.index_head_dim) - - # ========================================= - # Rotate activation - # ========================================= if self.config.dsa_indexer_rotate_activation: q = rotate_activation(q) - k = rotate_activation(k) # ========================================= - # Prepare weights for index scores + # k linear, k_norm, rotate, and weights_proj run in BF16 (FP8 disabled). # ========================================= - # [seqlen, batch, hidden_size] -> [seqlen, batch, index_n_heads] - weights, _ = self.linear_weights_proj(x) + with get_fp8_disabled_context(self.config): + # [seqlen, batch, hidden_size] -> [seqlen, batch, index_head_dim] + k, _ = self.linear_wk(x) + if self.config.dsa_indexer_k_norm_fp32: + k_dtype = k.dtype + k = self.k_norm(k.float()).to(dtype=k_dtype) + else: + k = self.k_norm(k) + # [seqlen, batch, index_head_dim] -> [seqlen, batch, 1, index_head_dim] + k = k.reshape(seqlen, bsz, 1, self.index_head_dim) + k = self._apply_rope(k, rotary_pos_emb, mscale, cu_seqlens=cu_seqlens_kv) + # [seqlen, batch, 1, index_head_dim] -> [seqlen, batch, index_head_dim] + k = k.reshape(seqlen, bsz, self.index_head_dim) + + # ========================================= + # Rotate activation (k only; q already rotated in FP8 path) + # ========================================= + if self.config.dsa_indexer_rotate_activation: + k = rotate_activation(k) + + # ========================================= + # Prepare weights for index scores + # ========================================= + # [seqlen, batch, hidden_size] -> [seqlen, batch, index_n_heads] + weights, _ = self.linear_weights_proj(x) weights = weights * (self.index_n_heads**-0.5) * self.softmax_scale return q, k, weights From 1631b38c8df6ec767835a597d46e82673f323d4d Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Fri, 4 Sep 2026 02:36:11 +0300 Subject: [PATCH 2/6] linear_wq_b also in bf16 Signed-off-by: Hollow Man --- .../core/transformer/experimental_attention_variant/dsa.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index b948a0662cd..1d606565d57 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -1433,10 +1433,11 @@ def forward_before_topk( seqlen, bsz, _ = x.size() # ========================================= - # q linear and apply rope to q (FP8) + # q linear and apply rope to q # ========================================= # [seqlen, batch, q_lora_rank] -> [seqlen, batch, index_n_heads * index_head_dim] - q, _ = self.linear_wq_b(qr) + with get_fp8_disabled_context(self.config): + q, _ = self.linear_wq_b(qr) # [seqlen, batch, index_n_heads * index_head_dim] # -> [seqlen, batch, index_n_heads, index_head_dim] q = q.reshape(seqlen, bsz, self.index_n_heads, self.index_head_dim) From 3dfe49f88af3f21732a833b375109db9ddc91fab Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Thu, 10 Sep 2026 00:49:48 +0300 Subject: [PATCH 3/6] KDA Signed-off-by: Hollow Man --- .../core/context_parallel_layout/__init__.py | 31 + .../context_parallel_layout/conversion.py | 668 ++++++++++++++ .../core/context_parallel_layout/routes.py | 280 ++++++ .../core/context_parallel_layout/types.py | 25 + .../core/context_parallel_layout/utils.py | 45 + .../core/extensions/transformer_engine.py | 1 + megatron/core/models/hybrid/hybrid_block.py | 12 + .../core/models/hybrid/hybrid_layer_specs.py | 51 ++ megatron/core/models/hybrid/hybrid_model.py | 24 +- megatron/core/models/hybrid/layers/utils.py | 3 + megatron/core/packed_seq_params.py | 16 + megatron/core/ssm/gated_delta_net/__init__.py | 8 + megatron/core/ssm/gated_delta_net/common.py | 39 +- megatron/core/ssm/gated_delta_net/kda.py | 813 ++++++++++++++++++ megatron/core/ssm/kda_layer_config.py | 10 + .../core/transformer/transformer_config.py | 21 + megatron/core/utils.py | 10 + .../unit_tests/ssm/test_kda_gate_precision.py | 115 +++ 18 files changed, 2158 insertions(+), 14 deletions(-) create mode 100644 megatron/core/context_parallel_layout/__init__.py create mode 100644 megatron/core/context_parallel_layout/conversion.py create mode 100644 megatron/core/context_parallel_layout/routes.py create mode 100644 megatron/core/context_parallel_layout/types.py create mode 100644 megatron/core/context_parallel_layout/utils.py create mode 100644 megatron/core/ssm/gated_delta_net/kda.py create mode 100644 megatron/core/ssm/kda_layer_config.py create mode 100644 tests/unit_tests/ssm/test_kda_gate_precision.py diff --git a/megatron/core/context_parallel_layout/__init__.py b/megatron/core/context_parallel_layout/__init__.py new file mode 100644 index 00000000000..06bb67e652c --- /dev/null +++ b/megatron/core/context_parallel_layout/__init__.py @@ -0,0 +1,31 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Public context parallel sequence partition-mode APIs. + +The implementation is split by responsibility; internal conversion and route-building +helpers remain in their respective submodules rather than being re-exported here. + +Ownership summary: + +- model builders choose the pipeline-stage input CP layout; +- blocks convert rank-local sequence tensors between layer preferences; +- model postprocess restores the public output boundary to the input layout; +- MTP validates its inner-layer layout preference but does not own outer conversion. +""" + +from megatron.core.context_parallel_layout.conversion import ( + CpPartitionModeConverter, + convert_module_input_tensors_cp_partition_mode, +) +from megatron.core.context_parallel_layout.routes import prebuild_thd_cp_partition_routes +from megatron.core.context_parallel_layout.types import CpPartitionMode, ThdCpRoute +from megatron.core.context_parallel_layout.utils import finalize_packed_seq_params + +__all__ = [ + "CpPartitionMode", + "CpPartitionModeConverter", + "ThdCpRoute", + "convert_module_input_tensors_cp_partition_mode", + "finalize_packed_seq_params", + "prebuild_thd_cp_partition_routes", +] diff --git a/megatron/core/context_parallel_layout/conversion.py b/megatron/core/context_parallel_layout/conversion.py new file mode 100644 index 00000000000..2763f2a8f06 --- /dev/null +++ b/megatron/core/context_parallel_layout/conversion.py @@ -0,0 +1,668 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Tensor operations for converting between CP partition modes.""" + +import warnings +from dataclasses import dataclass +from functools import lru_cache +from typing import TYPE_CHECKING, Any, Callable, Optional, Tuple, Union, cast + +import torch + +from megatron.core.context_parallel_layout.routes import ( + build_thd_cp_partition_route, + get_thd_cp_partition_route, +) +from megatron.core.context_parallel_layout.types import CpPartitionMode, ThdCpRoute +from megatron.core.context_parallel_layout.utils import ( + get_packed_seq_params_cp_partition_cu_seqlens, +) +from megatron.core.tensor_parallel.mappings import all_to_all +from megatron.core.utils import nvtx_range + +if TYPE_CHECKING: + from megatron.core.packed_seq_params import PackedSeqParams + + +class CpPartitionModeConverter: + """Convert tensors across one CP layout edge.""" + + def __init__( + self, + *, + packed_seq_params: Optional["PackedSeqParams"], + source_partition_mode: CpPartitionMode, + target_partition_mode: CpPartitionMode, + config: Any, + cp_group: Optional[torch.distributed.ProcessGroup] = None, + tp_group: Optional[torch.distributed.ProcessGroup] = None, + tp_cp_group: Optional[torch.distributed.ProcessGroup] = None, + ) -> None: + self.cp_group = cp_group + self.packed_seq_params = packed_seq_params + self.source_partition_mode = source_partition_mode + self.target_partition_mode = target_partition_mode + self.config = config + self.tp_group = tp_group + self.tp_cp_group = tp_cp_group + if ( + self.conversion_needed + and getattr(self.packed_seq_params, "qkv_format", None) == "thd" + and self.config.cuda_graph_impl == "full_iteration" + ): + raise ValueError( + "Full-iteration CUDA graph is not supported for THD CP layout conversion: " + f"source={self.source_partition_mode!r}, target={self.target_partition_mode!r}." + ) + + @property + def conversion_needed(self) -> bool: + """Return whether this edge needs a real layout conversion.""" + return ( + self.source_partition_mode != self.target_partition_mode + and self.cp_group is not None + and self.cp_group.size() > 1 + ) + + def assert_no_dense_attention_inputs( + self, + *, + attention_mask: Optional[torch.Tensor] = None, + attention_bias: Optional[torch.Tensor] = None, + hidden_states: Optional[torch.Tensor] = None, + ) -> None: + """Reject dense attention tensors when this edge would reorder tokens.""" + if not self.conversion_needed: + return + if attention_mask is not None: + self._raise_unsupported_dense_attention( + "an explicit attention_mask", hidden_states=hidden_states + ) + if attention_bias is not None: + self._raise_unsupported_dense_attention("attention_bias", hidden_states=hidden_states) + + def convert( + self, + value: Any, + *, + seq_dim: Union[int, Callable[[torch.Tensor], int]] = 0, + sequence_parallel: bool = False, + ) -> Any: + """Convert a tensor or nested tensor container across this layout edge.""" + if not self.conversion_needed or value is None: + return value + # Nested values may contain optional tensors; traverse containers while + # preserving their original shape. + if isinstance(value, tuple): + return tuple( + self.convert(part, seq_dim=seq_dim, sequence_parallel=sequence_parallel) + for part in value + ) + if isinstance(value, list): + return [ + self.convert(part, seq_dim=seq_dim, sequence_parallel=sequence_parallel) + for part in value + ] + if not torch.is_tensor(value): + return value + + resolved_seq_dim = seq_dim(value) if callable(seq_dim) else seq_dim + converted = convert_cp_partition_mode( + x=value, + source_partition_mode=self.source_partition_mode, + target_partition_mode=self.target_partition_mode, + seq_dim=resolved_seq_dim, + cu_seqlens=get_packed_seq_params_cp_partition_cu_seqlens(self.packed_seq_params), + sequence_parallel=sequence_parallel, + cp_group=self.cp_group, + tp_group=self.tp_group, + tp_cp_group=self.tp_cp_group, + thd_cp_partition_route=get_thd_cp_partition_route( + self.packed_seq_params, self.source_partition_mode, self.target_partition_mode + ), + ) + if self.packed_seq_params is not None: + self.packed_seq_params.cp_partition_mode = self.target_partition_mode + return converted + + def _raise_unsupported_dense_attention( + self, tensor_name: str, *, hidden_states: Optional[torch.Tensor] + ) -> None: + hidden_shape = tuple(hidden_states.shape) if hidden_states is not None else None + raise NotImplementedError( + "Changing CP partition mode with " + f"{tensor_name} is not supported yet: " + f"source={self.source_partition_mode!r}, " + f"target={self.target_partition_mode!r}, " + f"qkv_format={getattr(self.packed_seq_params, 'qkv_format', None)!r}, " + f"hidden_shape={hidden_shape}." + ) + + +def convert_module_input_tensors_cp_partition_mode( + *, + hidden_states: torch.Tensor, + packed_seq_params: Optional["PackedSeqParams"], + target_partition_mode: CpPartitionMode, + sequence_parallel: bool, + config: Any, + cp_group: Optional[torch.distributed.ProcessGroup] = None, + tp_group: Optional[torch.distributed.ProcessGroup] = None, + tp_cp_group: Optional[torch.distributed.ProcessGroup] = None, + attention_mask: Optional[torch.Tensor] = None, + attention_bias: Optional[torch.Tensor] = None, + key_value_states: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, Optional[CpPartitionModeConverter]]: + """Convert a module's rank-local sequence tensors to a target CP layout. + + This helper performs the common "entry conversion" pattern used by modules + that need to consume a different CP layout than their caller supplied. It + returns a converter for the opposite edge so the module output can be + converted back to the original input layout. + """ + if cp_group is None or cp_group.size() <= 1: + return hidden_states, None + + source_partition_mode = getattr(config, "cp_partition_mode", None) + if source_partition_mode is None: + raise ValueError( + "config.cp_partition_mode is required before module input CP layout conversion when " + "context parallelism is active." + ) + if source_partition_mode == target_partition_mode: + return hidden_states, None + + input_to_target_converter = CpPartitionModeConverter( + cp_group=cp_group, + packed_seq_params=packed_seq_params, + source_partition_mode=source_partition_mode, + target_partition_mode=target_partition_mode, + config=config, + tp_group=tp_group, + tp_cp_group=tp_cp_group, + ) + input_to_target_converter.assert_no_dense_attention_inputs( + attention_mask=attention_mask, attention_bias=attention_bias, hidden_states=hidden_states + ) + if key_value_states is not None: + raise NotImplementedError( + "Changing CP partition mode with cross-attention key/value states is not supported " + f"yet: source={source_partition_mode!r}, target={target_partition_mode!r}." + ) + hidden_states = input_to_target_converter.convert( + value=hidden_states, seq_dim=0, sequence_parallel=sequence_parallel + ) + + target_to_input_converter = CpPartitionModeConverter( + cp_group=cp_group, + packed_seq_params=packed_seq_params, + source_partition_mode=target_partition_mode, + target_partition_mode=source_partition_mode, + config=config, + tp_group=tp_group, + tp_cp_group=tp_cp_group, + ) + return (hidden_states, target_to_input_converter) + + +def convert_cp_partition_mode( + x: torch.Tensor, + *, + source_partition_mode: Optional[str], + target_partition_mode: Optional[str], + seq_dim: int = 0, + cu_seqlens: Optional[torch.Tensor] = None, + sequence_parallel: bool = False, + cp_group: Optional[torch.distributed.ProcessGroup] = None, + tp_group: Optional[torch.distributed.ProcessGroup] = None, + tp_cp_group: Optional[torch.distributed.ProcessGroup] = None, + thd_cp_partition_route: Optional[ThdCpRoute] = None, +) -> torch.Tensor: + """Convert a sequence tensor between CP zigzag and contiguous layouts. + + SBHD tensors use one unified all-to-all-v redistribution path over CP or + TPxCP. THD tensors use their packed-token CP route and, when sequence + parallelism shards the packed sequence, retain the naive TP gather/scatter + fallback. + """ + + if source_partition_mode == target_partition_mode: + return x + + cp_size = cp_group.size() if cp_group is not None else 1 + if cp_size == 1: + return x + assert cp_group is not None + + if source_partition_mode not in ("zigzag", "contiguous") or target_partition_mode not in ( + "zigzag", + "contiguous", + ): + cp_rank = cp_group.rank() if cp_group is not None else 0 + raise ValueError( + f"Unsupported CP partition mode conversion " + f"{source_partition_mode!r} -> {target_partition_mode!r}; " + f"shape={tuple(x.shape)}, seq_dim={seq_dim}, cp_size={cp_size}, cp_rank={cp_rank}." + ) + source_layout = cast(CpPartitionMode, source_partition_mode) + target_layout = cast(CpPartitionMode, target_partition_mode) + + if cu_seqlens is None: + moved = x.movedim(seq_dim, 0) if seq_dim != 0 else x + converted = _redistribute_sbhd_layout( + input_=moved, + cp_group=cp_group, + source_layout=source_layout, + target_layout=target_layout, + sequence_parallel=sequence_parallel, + tp_group=tp_group, + tp_cp_group=tp_cp_group, + ) + return converted.movedim(0, seq_dim).contiguous() if seq_dim != 0 else converted + + if sequence_parallel and tp_group is not None and tp_group.size() > 1: + from megatron.core.tensor_parallel.mappings import ( + gather_from_sequence_parallel_region, + scatter_to_sequence_parallel_region, + ) + + # TODO(yuzhongw): replace the naive THD TP gather -> CP all-to-all -> TP scatter + # fallback with a direct packed THD TPxCP redistribution path. + warnings.warn( + "THD CP layout conversion with sequence parallelism uses the naive " + "TP gather -> CP all-to-all -> TP scatter fallback.", + RuntimeWarning, + stacklevel=2, + ) + moved = x.movedim(seq_dim, 0) if seq_dim != 0 else x + # This gather is only used to run a duplicated CP layout permutation before + # scattering back to SP shards. Its backward must split, not reduce-scatter; + # otherwise every TP rank contributes the same full-sequence gradient. + gathered = gather_from_sequence_parallel_region( + input_=moved, tensor_parallel_output_grad=False, group=tp_group + ) + converted = _redistribute_thd_layout( + x=gathered, + cp_group=cp_group, + seq_dim=0, + cu_seqlens=cu_seqlens, + source_partition_mode=source_layout, + target_partition_mode=target_layout, + thd_cp_partition_route=thd_cp_partition_route, + ) + scattered = scatter_to_sequence_parallel_region(input_=converted, group=tp_group) + return scattered.movedim(0, seq_dim).contiguous() if seq_dim != 0 else scattered + + return _redistribute_thd_layout( + x=x, + cp_group=cp_group, + seq_dim=seq_dim, + cu_seqlens=cu_seqlens, + source_partition_mode=source_layout, + target_partition_mode=target_layout, + thd_cp_partition_route=thd_cp_partition_route, + ) + + +def _pack_thd_cp_route_send_buffer( + x: torch.Tensor, send_index: Optional[torch.Tensor] +) -> torch.Tensor: + if send_index is None: + return x + return x.index_select(0, send_index) + + +def _scatter_thd_cp_route_recv_buffer( + recv_buf: torch.Tensor, recv_index: Optional[torch.Tensor], out_shape: Tuple[int, ...] +) -> torch.Tensor: + if recv_index is None: + return recv_buf + out = recv_buf.new_empty(out_shape) + if recv_index.numel() > 0: + out.index_copy_(0, recv_index, recv_buf) + return out + + +def _redistribute_thd_layout( + x: torch.Tensor, + cp_group: Optional[torch.distributed.ProcessGroup], + seq_dim: int, + cu_seqlens: torch.Tensor, + source_partition_mode: str, + target_partition_mode: str, + thd_cp_partition_route: Optional[ThdCpRoute] = None, +) -> torch.Tensor: + """Single-all-to-all THD permutation between zigzag and contiguous layouts. + + The packed THD tensor stays packed: we first group local tokens by their + target CP rank, exchange those groups once, then scatter received tokens + back into the target rank-local order. + """ + cp_size = cp_group.size() if cp_group is not None else 1 + if cp_size == 1: + return x + assert cp_group is not None + cp_rank = cp_group.rank() + conversion_name = f"{source_partition_mode}_to_{target_partition_mode}" + with nvtx_range(f"cp_layout/thd/swap/{conversion_name}"): + if seq_dim != 0: + x = x.movedim(seq_dim, 0) + x = x.contiguous() + + route = thd_cp_partition_route + if route is None: + route = build_thd_cp_partition_route( + cu_seqlens=cu_seqlens, cp_size=cp_size, cp_rank=cp_rank, device=x.device + ) + + if source_partition_mode == "zigzag" and target_partition_mode == "contiguous": + send_index = route.zigzag_index + recv_index = route.contiguous_index + input_split_sizes = route.zigzag_split_sizes + output_split_sizes = route.contiguous_split_sizes + elif source_partition_mode == "contiguous" and target_partition_mode == "zigzag": + send_index = route.contiguous_index + recv_index = route.zigzag_index + input_split_sizes = route.contiguous_split_sizes + output_split_sizes = route.zigzag_split_sizes + else: + raise ValueError( + f"Unsupported CP partition mode conversion " + f"{source_partition_mode!r} -> {target_partition_mode!r} for THD route." + ) + + local_source_length = sum(input_split_sizes) + local_target_length = sum(output_split_sizes) + + if x.size(0) != local_source_length: + raise ValueError( + f"Local THD tensor length ({x.size(0)}) does not match {source_partition_mode} " + f"rank-{cp_rank} partition length ({local_source_length})." + ) + if local_target_length != x.size(0): + raise ValueError( + "THD CP layout conversion must preserve the local token count, " + f"got source={local_source_length}, target={local_target_length}, " + f"cp_size={cp_size}, cp_rank={cp_rank}, " + f"source_layout={source_partition_mode!r}, " + f"target_layout={target_partition_mode!r}." + ) + + with nvtx_range(f"cp_layout/thd/pack/{conversion_name}"): + send_buf = _pack_thd_cp_route_send_buffer(x=x, send_index=send_index) + if not send_buf.is_contiguous(): + send_buf = send_buf.contiguous() + + with nvtx_range(f"cp_layout/thd/all_to_all/{conversion_name}"): + recv_buf = all_to_all( + group=cp_group, + input_=send_buf, + output_split_sizes_=output_split_sizes, + input_split_sizes=input_split_sizes, + ) + + with nvtx_range(f"cp_layout/thd/scatter/{conversion_name}"): + out_shape = (local_target_length,) + tuple(x.shape[1:]) + out = _scatter_thd_cp_route_recv_buffer( + recv_buf=recv_buf, recv_index=recv_index, out_shape=out_shape + ) + + if seq_dim != 0: + out = out.movedim(0, seq_dim) + return out.contiguous() + + +@dataclass(frozen=True) +class _SbhdLayoutRedistributionPlan: + """Rank-local SBHD all-to-all plan expressed in sequence-segment counts.""" + + send_slots: tuple[int, ...] + input_segment_counts: tuple[int, ...] + output_segment_counts: tuple[int, ...] + receive_permutation: tuple[int, ...] + + +def _sbhd_segments_per_rank(tp_size: int) -> int: + """Return two SBHD segments for CP-only conversion and one for even-TP SP conversion.""" + if tp_size == 1: + return 2 + if tp_size % 2 != 0: + raise ValueError( + "Sequence-parallel SBHD CP layout conversion requires an even tensor-parallel size, " + f"got {tp_size}" + ) + return 1 + + +def _local_sbhd_segment_ids( + layout: CpPartitionMode, cp_size: int, cp_rank: int, tp_size: int = 1, tp_rank: int = 0 +) -> tuple[int, ...]: + """Return the atomic SBHD sequence segments owned by one TP×CP rank.""" + segments_per_rank = _sbhd_segments_per_rank(tp_size=tp_size) + if layout == "contiguous": + first_segment = segments_per_rank * (cp_rank * tp_size + tp_rank) + return tuple(range(first_segment, first_segment + segments_per_rank)) + if layout == "zigzag": + segments_per_cp_half = tp_size * segments_per_rank // 2 + front_start = cp_rank * segments_per_cp_half + back_start = (2 * cp_size - cp_rank - 1) * segments_per_cp_half + cp_segments = tuple(range(front_start, front_start + segments_per_cp_half)) + tuple( + range(back_start, back_start + segments_per_cp_half) + ) + sp_start = segments_per_rank * tp_rank + return cp_segments[sp_start : sp_start + segments_per_rank] + raise ValueError(f"Unsupported CP layout: {layout}") + + +@lru_cache(maxsize=None) +def _sbhd_segment_owner( + segment_id: int, layout: CpPartitionMode, cp_size: int, tp_size: int +) -> tuple[int, int]: + for cp_rank in range(cp_size): + for tp_rank in range(tp_size): + if segment_id in _local_sbhd_segment_ids( + layout=layout, cp_size=cp_size, cp_rank=cp_rank, tp_size=tp_size, tp_rank=tp_rank + ): + return cp_rank, tp_rank + raise ValueError( + f"SBHD segment {segment_id} is not present in the {layout} layout for " + f"{cp_size=} and {tp_size=}" + ) + + +@lru_cache(maxsize=None) +def _build_sbhd_group_rank_by_logical_rank( + cp_global_ranks: tuple[int, ...], + tp_global_ranks: tuple[int, ...], + tp_cp_global_ranks: tuple[int, ...], + current_global_rank: int, +) -> tuple[int, ...]: + """Map logical ``cp_rank * tp_size + tp_rank`` coordinates to group ranks for SBHD.""" + group_rank_by_global_rank = { + global_rank: group_rank for group_rank, global_rank in enumerate(tp_cp_global_ranks) + } + group_rank_by_logical_rank = [] + for cp_global_rank in cp_global_ranks: + for tp_global_rank in tp_global_ranks: + target_global_rank = cp_global_rank + tp_global_rank - current_global_rank + if target_global_rank not in group_rank_by_global_rank: + raise RuntimeError( + "TP and CP process groups do not form the expected Cartesian product" + ) + group_rank_by_logical_rank.append(group_rank_by_global_rank[target_global_rank]) + return tuple(group_rank_by_logical_rank) + + +def _get_sbhd_group_rank_by_logical_rank( + cp_group: torch.distributed.ProcessGroup, + tp_group: torch.distributed.ProcessGroup, + tp_cp_group: torch.distributed.ProcessGroup, +) -> tuple[int, ...]: + return _build_sbhd_group_rank_by_logical_rank( + cp_global_ranks=tuple(torch.distributed.get_process_group_ranks(cp_group)), + tp_global_ranks=tuple(torch.distributed.get_process_group_ranks(tp_group)), + tp_cp_global_ranks=tuple(torch.distributed.get_process_group_ranks(tp_cp_group)), + current_global_rank=torch.distributed.get_rank(), + ) + + +@lru_cache(maxsize=None) +def _build_sbhd_layout_redistribution_plan( + source_layout: CpPartitionMode, + target_layout: CpPartitionMode, + cp_size: int, + cp_rank: int, + tp_size: int = 1, + tp_rank: int = 0, + group_rank_by_logical_rank: tuple[int, ...] | None = None, +) -> _SbhdLayoutRedistributionPlan: + """Build the SBHD all-to-all-v plan for one rank of a CP layout conversion.""" + if cp_size < 1: + raise ValueError(f"cp_size must be positive, got {cp_size}") + if tp_size < 1: + raise ValueError(f"tp_size must be positive, got {tp_size}") + if not 0 <= cp_rank < cp_size: + raise ValueError(f"cp_rank must be in [0, {cp_size}), got {cp_rank}") + if not 0 <= tp_rank < tp_size: + raise ValueError(f"tp_rank must be in [0, {tp_size}), got {tp_rank}") + + group_size = cp_size * tp_size + if group_rank_by_logical_rank is None: + group_rank_by_logical_rank = tuple(range(group_size)) + if sorted(group_rank_by_logical_rank) != list(range(group_size)): + raise ValueError("group_rank_by_logical_rank must be a permutation of the group ranks") + + source_ids = _local_sbhd_segment_ids( + layout=source_layout, cp_size=cp_size, cp_rank=cp_rank, tp_size=tp_size, tp_rank=tp_rank + ) + target_ids = _local_sbhd_segment_ids( + layout=target_layout, cp_size=cp_size, cp_rank=cp_rank, tp_size=tp_size, tp_rank=tp_rank + ) + + def destination_group_rank(segment_id: int) -> int: + destination_cp_rank, destination_tp_rank = _sbhd_segment_owner( + segment_id=segment_id, layout=target_layout, cp_size=cp_size, tp_size=tp_size + ) + destination_logical_rank = destination_cp_rank * tp_size + destination_tp_rank + return group_rank_by_logical_rank[destination_logical_rank] + + send_entries = sorted( + (destination_group_rank(segment_id=segment_id), slot) + for slot, segment_id in enumerate(source_ids) + ) + send_slots = tuple(slot for _, slot in send_entries) + input_segment_counts = tuple( + sum(destination == rank for destination, _ in send_entries) for rank in range(group_size) + ) + + received_ids = [] + output_segment_counts = [] + source_logical_ranks = sorted( + range(group_size), key=lambda logical_rank: group_rank_by_logical_rank[logical_rank] + ) + for source_logical_rank in source_logical_ranks: + source_cp_rank, source_tp_rank = divmod(source_logical_rank, tp_size) + rank_source_ids = _local_sbhd_segment_ids( + layout=source_layout, + cp_size=cp_size, + cp_rank=source_cp_rank, + tp_size=tp_size, + tp_rank=source_tp_rank, + ) + ids_from_source = [ + segment_id + for segment_id in rank_source_ids + if _sbhd_segment_owner( + segment_id=segment_id, layout=target_layout, cp_size=cp_size, tp_size=tp_size + ) + == (cp_rank, tp_rank) + ] + received_ids.extend(ids_from_source) + output_segment_counts.append(len(ids_from_source)) + + if sorted(received_ids) != sorted(target_ids): + raise RuntimeError( + f"Invalid {source_layout}-to-{target_layout} SBHD redistribution plan for " + f"CP rank {cp_rank}, TP rank {tp_rank}: received {received_ids}, " + f"expected {target_ids}" + ) + receive_permutation = tuple(received_ids.index(segment_id) for segment_id in target_ids) + + return _SbhdLayoutRedistributionPlan( + send_slots=send_slots, + input_segment_counts=input_segment_counts, + output_segment_counts=tuple(output_segment_counts), + receive_permutation=receive_permutation, + ) + + +def _redistribute_sbhd_layout( + input_: torch.Tensor, + cp_group: torch.distributed.ProcessGroup, + source_layout: CpPartitionMode, + target_layout: CpPartitionMode, + sequence_parallel: bool, + tp_group: Optional[torch.distributed.ProcessGroup], + tp_cp_group: Optional[torch.distributed.ProcessGroup], +) -> torch.Tensor: + """Redistribute local SBHD sequence segments with a differentiable all-to-all-v.""" + cp_size = cp_group.size() + if cp_size == 1 or source_layout == target_layout: + return input_ + + cp_rank = cp_group.rank() + tp_size, tp_rank = 1, 0 + communication_group = cp_group + group_rank_by_logical_rank = None + if sequence_parallel and tp_group is not None and tp_group.size() > 1: + if tp_cp_group is None: + raise ValueError( + "tp_cp_group is required for direct sequence-parallel SBHD layout conversion" + ) + tp_size, tp_rank = tp_group.size(), tp_group.rank() + communication_group = tp_cp_group + group_rank_by_logical_rank = _get_sbhd_group_rank_by_logical_rank( + cp_group=cp_group, tp_group=tp_group, tp_cp_group=tp_cp_group + ) + + plan = _build_sbhd_layout_redistribution_plan( + source_layout=source_layout, + target_layout=target_layout, + cp_size=cp_size, + cp_rank=cp_rank, + tp_size=tp_size, + tp_rank=tp_rank, + group_rank_by_logical_rank=group_rank_by_logical_rank, + ) + + input_contiguous = input_.contiguous() + local_seq_len = input_contiguous.shape[0] + local_segment_count = _sbhd_segments_per_rank(tp_size=tp_size) + if local_seq_len % local_segment_count != 0: + raise ValueError( + "SBHD CP layout conversion requires the sequence length local to each TP×CP rank to " + f"be divisible by {local_segment_count}, got {local_seq_len}" + ) + segment_len = local_seq_len // local_segment_count + segment_shape = (local_segment_count, segment_len, *input_contiguous.shape[1:]) + segments = input_contiguous.reshape(segment_shape) + + if plan.send_slots == tuple(range(local_segment_count)): + send_buffer = input_contiguous + else: + send_buffer = segments.flip(0).reshape(input_contiguous.shape) + input_split_sizes = [count * segment_len for count in plan.input_segment_counts] + output_split_sizes = [count * segment_len for count in plan.output_segment_counts] + received = all_to_all( + group=communication_group, + input_=send_buffer, + output_split_sizes_=output_split_sizes, + input_split_sizes=input_split_sizes, + ) + + received_segments = received.reshape(segment_shape) + if plan.receive_permutation == tuple(range(local_segment_count)): + output = received + else: + output = received_segments.flip(0).reshape(input_contiguous.shape) + return output.contiguous() diff --git a/megatron/core/context_parallel_layout/routes.py b/megatron/core/context_parallel_layout/routes.py new file mode 100644 index 00000000000..66981dfca52 --- /dev/null +++ b/megatron/core/context_parallel_layout/routes.py @@ -0,0 +1,280 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""THD context-parallel route helpers.""" + +import warnings +from typing import TYPE_CHECKING, List, Optional, Tuple + +import torch + +from megatron.core.context_parallel_layout.types import CpPartitionMode, ThdCpRoute +from megatron.core.context_parallel_layout.utils import ( + get_packed_seq_params_cp_partition_cu_seqlens, +) +from megatron.core.utils import nvtx_range + +if TYPE_CHECKING: + from megatron.core.packed_seq_params import PackedSeqParams + +_ThdLayoutSegment = Tuple[int, int, int] + + +def _compact_thd_cu_seqlens_to_list(cu_seqlens: torch.Tensor) -> List[int]: + if cu_seqlens.dim() != 1: + raise ValueError(f"cu_seqlens must be 1-D, got shape {tuple(cu_seqlens.shape)}.") + + cu = cu_seqlens.detach().to(device="cpu", dtype=torch.long).tolist() + if not cu or cu[0] != 0: + raise ValueError(f"cu_seqlens must start at 0, got {cu_seqlens}.") + + compact_cu: List[int] = [cu[0]] + prev = cu[0] + for value in cu[1:]: + if value < prev: + raise ValueError(f"cu_seqlens must be nondecreasing, got {cu_seqlens}.") + if value != prev: + compact_cu.append(value) + prev = value + return compact_cu + + +def _validate_thd_route_partitioning(cu: List[int], cp_size: int) -> None: + total_tokens = cu[-1] + if total_tokens % cp_size != 0: + raise ValueError( + f"Contiguous CP partitioning requires total_tokens={total_tokens} " + f"to be divisible by cp_size={cp_size}." + ) + + chunk_divisor = 2 * cp_size + bad_seq_lens = [ + seq_end - seq_start + for seq_start, seq_end in zip(cu[:-1], cu[1:]) + if (seq_end - seq_start) % chunk_divisor != 0 + ] + if bad_seq_lens: + raise ValueError( + "All packed sequence lengths must be divisible by " + f"2 * cp_size ({chunk_divisor}) for zigzag CP layout conversion, " + f"got {bad_seq_lens}." + ) + + +def _build_thd_layout_segments( + cu: List[int], cp_size: int, cp_rank: int, cp_partition_mode: CpPartitionMode +) -> Tuple[List[_ThdLayoutSegment], int]: + total_tokens = cu[-1] + if cp_partition_mode == "contiguous": + part_len = total_tokens // cp_size + if part_len == 0: + return [], 0 + return [(cp_rank * part_len, part_len, 0)], part_len + + if cp_partition_mode != "zigzag": + raise ValueError( + f"Unsupported context-parallel partition mode {cp_partition_mode!r} " + f"for THD layout segments with cp_size={cp_size}, rank={cp_rank}." + ) + + segments: List[_ThdLayoutSegment] = [] + local_start = 0 + for seq_start, seq_end in zip(cu[:-1], cu[1:]): + seq_len = seq_end - seq_start + chunk_len = seq_len // (2 * cp_size) + first_chunk = cp_rank + second_chunk = 2 * cp_size - cp_rank - 1 + segments.append((seq_start + first_chunk * chunk_len, chunk_len, local_start)) + segments.append((seq_start + second_chunk * chunk_len, chunk_len, local_start + chunk_len)) + local_start += 2 * chunk_len + + return segments, local_start + + +def _intersect_thd_layout_segments( + source_segments: List[_ThdLayoutSegment], target_segments: List[_ThdLayoutSegment] +) -> List[Tuple[int, int, int]]: + intersections: List[Tuple[int, int, int]] = [] + source_index = 0 + target_index = 0 + while source_index < len(source_segments) and target_index < len(target_segments): + source_global_start, source_len, source_local_start = source_segments[source_index] + target_global_start, target_len, target_local_start = target_segments[target_index] + source_global_end = source_global_start + source_len + target_global_end = target_global_start + target_len + + overlap_start = max(source_global_start, target_global_start) + overlap_end = min(source_global_end, target_global_end) + if overlap_start < overlap_end: + intersections.append( + ( + source_local_start + overlap_start - source_global_start, + target_local_start + overlap_start - target_global_start, + overlap_end - overlap_start, + ) + ) + + if source_global_end <= target_global_end: + source_index += 1 + else: + target_index += 1 + + return intersections + + +def _build_thd_layout_side_route( + local_segments: List[_ThdLayoutSegment], + target_segments_by_rank: List[List[_ThdLayoutSegment]], + *, + device: torch.device, +) -> Tuple[Optional[torch.Tensor], List[int]]: + row_order: List[int] = [] + split_sizes: List[int] = [] + for peer_rank in range(len(target_segments_by_rank)): + intersections = _intersect_thd_layout_segments( + local_segments, target_segments_by_rank[peer_rank] + ) + intersections.sort(key=lambda item: item[1]) + split_size = 0 + for source_row, _, length in intersections: + row_order.extend(range(source_row, source_row + length)) + split_size += length + split_sizes.append(split_size) + + if all(row == index for index, row in enumerate(row_order)): + return None, split_sizes + return torch.tensor(row_order, device=device, dtype=torch.long), split_sizes + + +def build_thd_cp_partition_route( + cu_seqlens: torch.Tensor, cp_size: int, cp_rank: int, *, device: Optional[torch.device] = None +) -> ThdCpRoute: + """Precompute the rank-local THD CP layout route for a microbatch. + + The route stores both zigzag and contiguous layout views and can be reused + for either conversion direction over tensors with the same THD sequence + axis in the same microbatch. + """ + if cp_size < 1: + raise ValueError(f"cp_size must be >= 1, got {cp_size}.") + if not 0 <= cp_rank < cp_size: + raise ValueError(f"cp_rank must be in [0, {cp_size}), got {cp_rank}.") + if device is None: + device = cu_seqlens.device + + with nvtx_range("cp_layout/thd/route"): + cu = _compact_thd_cu_seqlens_to_list(cu_seqlens) + _validate_thd_route_partitioning(cu, cp_size) + + zigzag_segments_by_rank: List[List[_ThdLayoutSegment]] = [] + zigzag_lengths: List[int] = [] + contiguous_segments_by_rank: List[List[_ThdLayoutSegment]] = [] + for rank in range(cp_size): + zigzag_segments, zigzag_length = _build_thd_layout_segments(cu, cp_size, rank, "zigzag") + contiguous_segments, contiguous_length = _build_thd_layout_segments( + cu, cp_size, rank, "contiguous" + ) + if zigzag_length != contiguous_length: + raise ValueError( + "THD CP layout conversion must preserve local token count, " + f"got zigzag={zigzag_length}, contiguous={contiguous_length} " + f"for cp_size={cp_size}, rank={rank}." + ) + zigzag_segments_by_rank.append(zigzag_segments) + zigzag_lengths.append(zigzag_length) + contiguous_segments_by_rank.append(contiguous_segments) + + zigzag_index, zigzag_split_sizes = _build_thd_layout_side_route( + zigzag_segments_by_rank[cp_rank], contiguous_segments_by_rank, device=device + ) + contiguous_index, contiguous_split_sizes = _build_thd_layout_side_route( + contiguous_segments_by_rank[cp_rank], zigzag_segments_by_rank, device=device + ) + + local_length = zigzag_lengths[cp_rank] + if sum(zigzag_split_sizes) != local_length: + raise ValueError( + "Zigzag THD CP route split sizes do not match the local token count: " + f"splits={zigzag_split_sizes}, local_length={local_length}." + ) + if sum(contiguous_split_sizes) != local_length: + raise ValueError( + "Contiguous THD CP route split sizes do not match the local token count: " + f"splits={contiguous_split_sizes}, local_length={local_length}." + ) + + return ThdCpRoute( + zigzag_index=zigzag_index, + zigzag_split_sizes=zigzag_split_sizes, + contiguous_index=contiguous_index, + contiguous_split_sizes=contiguous_split_sizes, + ) + + +def get_thd_cp_partition_route( + packed_seq_params: Optional["PackedSeqParams"], + source_partition_mode: CpPartitionMode, + target_partition_mode: CpPartitionMode, +) -> Optional[ThdCpRoute]: + """Return the precomputed THD CP partition route for one direction. + + The fallback below is intentionally only a compatibility path: it performs + a blocking device-to-host copy while compacting ``cu_seqlens`` and mutates + ``packed_seq_params`` by storing the resulting route. Production callers + should prebuild routes when constructing the batch. + """ + if source_partition_mode == target_partition_mode: + return None + if source_partition_mode not in ("zigzag", "contiguous") or target_partition_mode not in ( + "zigzag", + "contiguous", + ): + raise ValueError( + f"Unsupported CP partition mode conversion " + f"{source_partition_mode!r} -> {target_partition_mode!r} for THD route." + ) + if packed_seq_params is None or getattr(packed_seq_params, "qkv_format", None) != "thd": + return None + + route = getattr(packed_seq_params, "cp_partition_route", None) + if route is not None: + return route + + warnings.warn( + "THD PackedSeqParams is missing precomputed context-parallel layout routes. " + "This lookup will attempt to build them from packed_seq_params.cp_group as " + "a compatibility fallback. The fallback synchronizes cu_seqlens to CPU " + "and mutates packed_seq_params.cp_partition_route, so it should not be " + "used on the steady-state forward path. Callers should prebuild THD CP " + "routes when constructing the batch; a future release will require the " + "routes to be present before layout conversion.", + FutureWarning, + stacklevel=2, + ) + prebuild_thd_cp_partition_routes(packed_seq_params) + return getattr(packed_seq_params, "cp_partition_route", None) + + +def prebuild_thd_cp_partition_routes( + packed_seq_params: Optional["PackedSeqParams"], + cp_group: Optional[torch.distributed.ProcessGroup] = None, + *, + device: Optional[torch.device] = None, +) -> None: + """Prebuild the THD CP layout route for a packed microbatch.""" + if packed_seq_params is None or getattr(packed_seq_params, "qkv_format", None) != "thd": + return + if cp_group is None: + cp_group = getattr(packed_seq_params, "cp_group", None) + if cp_group is None or cp_group.size() <= 1: + return + cp_size = cp_group.size() + cp_rank = cp_group.rank() + cu_seqlens = get_packed_seq_params_cp_partition_cu_seqlens(packed_seq_params) + if cu_seqlens is None: + return + if device is None: + device = cu_seqlens.device + + packed_seq_params.cp_partition_route = build_thd_cp_partition_route( + cu_seqlens, cp_size, cp_rank, device=device + ) diff --git a/megatron/core/context_parallel_layout/types.py b/megatron/core/context_parallel_layout/types.py new file mode 100644 index 00000000000..97a54bef25a --- /dev/null +++ b/megatron/core/context_parallel_layout/types.py @@ -0,0 +1,25 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Leaf type definitions for context-parallel layout helpers.""" + +from dataclasses import dataclass +from typing import List, Literal, Optional + +import torch + +CpPartitionMode = Literal["zigzag", "contiguous"] + + +@dataclass +class ThdCpRoute: + """Rank-local route plan for THD zigzag/contiguous CP layout conversion. + + The route stores each layout's local communication view exactly once. A + directional conversion interprets the source layout fields as send metadata + and the target layout fields as receive metadata. + """ + + zigzag_index: Optional[torch.Tensor] + zigzag_split_sizes: List[int] + contiguous_index: Optional[torch.Tensor] + contiguous_split_sizes: List[int] diff --git a/megatron/core/context_parallel_layout/utils.py b/megatron/core/context_parallel_layout/utils.py new file mode 100644 index 00000000000..66d9b9183ae --- /dev/null +++ b/megatron/core/context_parallel_layout/utils.py @@ -0,0 +1,45 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Packed-sequence metadata helpers for CP partition-mode tracking.""" + +from typing import TYPE_CHECKING, Optional + +import torch + +if TYPE_CHECKING: + from megatron.core.packed_seq_params import PackedSeqParams + + +def get_packed_seq_params_cp_partition_cu_seqlens( + packed_seq_params: Optional["PackedSeqParams"], +) -> Optional[torch.Tensor]: + """Return THD cumulative sequence lengths used for CP layout conversion. + + ``packed_seq_params=None`` represents the ordinary SBHD path. Only THD + metadata carries global packed-token boundaries. + """ + if packed_seq_params is None or getattr(packed_seq_params, "qkv_format", None) != "thd": + return None + return ( + packed_seq_params.cu_seqlens_q_padded + if packed_seq_params.cu_seqlens_q_padded is not None + else packed_seq_params.cu_seqlens_q + ) + + +def finalize_packed_seq_params( + packed_seq_params: Optional["PackedSeqParams"], +) -> Optional["PackedSeqParams"]: + """Resolve CP metadata and prebuild the THD layout route for a microbatch.""" + if packed_seq_params is None: + return None + + # Keep these imports local: routes depends on this module for metadata access. + from megatron.core.context_parallel_layout.routes import prebuild_thd_cp_partition_routes + from megatron.core.packed_seq_params import resolve_cp_group + from megatron.core.parallel_state import get_context_parallel_group + + cp_group = resolve_cp_group(get_context_parallel_group(), packed_seq_params) + packed_seq_params.cp_group = cp_group + prebuild_thd_cp_partition_routes(packed_seq_params, cp_group) + return packed_seq_params diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 53203cdf428..207a7721610 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -2251,6 +2251,7 @@ def __init__( self.kept_packed_seq_params.discard("seq_idx") self.kept_packed_seq_params.discard("tokens_per_sample") self.kept_packed_seq_params.discard("cp_scatter_cache") + self.kept_packed_seq_params.discard("cp_partition_mode") if get_te_version() < PkgVersion("2.2.0"): self.kept_packed_seq_params.discard("pad_between_seqs") diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py index 9160fb2ee7b..6a69452d345 100644 --- a/megatron/core/models/hybrid/hybrid_block.py +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -61,6 +61,7 @@ class HybridStackSubmodules: mamba_layer: Union[ModuleSpec, type] = IdentityOp gdn_layer: Union[ModuleSpec, type] = IdentityOp + kda_layer: Union[ModuleSpec, type] = IdentityOp attention_layer: Union[ModuleSpec, type] = IdentityOp dsa_layer: Union[ModuleSpec, type] = IdentityOp mla_layer: Union[ModuleSpec, type] = IdentityOp @@ -286,6 +287,17 @@ def __init__( pp_layer_offset=pp_layer_offset, name=(name + f".layers.{i}") if name is not None else None, ) + elif type(layer_config) is layer_utils.KDALayerConfig: + layer = build_module( + submodules.kda_layer, + config=layer_config, + layer_number=layer_number, + pg_collection=pg_collection, + # Set to False as we do not want to change offset. + add_layer_offset=False, + pp_layer_offset=pp_layer_offset, + name=(name + f".layers.{i}") if name is not None else None, + ) else: raise ValueError( f"Unexpected hybrid layer config type: {type(layer_config).__name__}" diff --git a/megatron/core/models/hybrid/hybrid_layer_specs.py b/megatron/core/models/hybrid/hybrid_layer_specs.py index f0be7ef891b..1c6b2b28a30 100755 --- a/megatron/core/models/hybrid/hybrid_layer_specs.py +++ b/megatron/core/models/hybrid/hybrid_layer_specs.py @@ -16,6 +16,7 @@ ) from megatron.core.models.hybrid.hybrid_block import HybridStack, HybridStackSubmodules from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNetSubmodules +from megatron.core.ssm.gated_delta_net.kda import KimiDeltaAttention, KimiDeltaAttentionSubmodules from megatron.core.ssm.gated_delta_product import ( GatedDeltaProductMixer, GatedDeltaProductMixerSubmodules, @@ -137,6 +138,30 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): self_attn_bda=get_bias_dropout_add, ), ), + kda_layer=ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + input_layernorm=TENorm, + self_attention=ModuleSpec( + module=KimiDeltaAttention, + submodules=KimiDeltaAttentionSubmodules( + in_proj=TEColumnParallelLinear, + beta_proj=TEColumnParallelLinear, + # Two-stage low-rank gates (GLM-5.3-Flash); only used when + # config.kda_two_stage_gates=True, else IdentityOp. + # f_a/g_a are replicated (TELinear parallel_mode="duplicated"); + # f_b/g_b are TP-sharded column-parallel. + f_a_proj=TELinear, + f_b_proj=TEColumnParallelLinear, + g_a_proj=TELinear, + g_b_proj=TEColumnParallelLinear, + out_norm=TENorm, + out_proj=TERowParallelLinear, + ), + ), + self_attn_bda=get_bias_dropout_add, + ), + ), # Started with spec from gpt_layer_specs.py (with MLP removed) # Using the TE spec because we had problems getting the non-TE spec # working @@ -245,6 +270,7 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): TELayerNormColumnParallelLinear, TERowParallelLinear ), gdn_layer=hybrid_stack_spec.submodules.gdn_layer, + kda_layer=hybrid_stack_spec.submodules.kda_layer, attention_layer=hybrid_stack_spec.submodules.attention_layer, dsa_layer=hybrid_stack_spec.submodules.dsa_layer, mlp_layer=hybrid_stack_spec.submodules.mlp_layer, @@ -284,6 +310,30 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): self_attn_bda=get_bias_dropout_add, ), ), + kda_layer=ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + input_layernorm=TENorm, + self_attention=ModuleSpec( + module=KimiDeltaAttention, + submodules=KimiDeltaAttentionSubmodules( + in_proj=InferenceColumnParallelLinear, + beta_proj=InferenceColumnParallelLinear, + # Two-stage low-rank gates (GLM-5.3-Flash); only used when + # config.kda_two_stage_gates=True, else IdentityOp. + # f_a/g_a replicated (TELinear parallel_mode="duplicated"); + # f_b/g_b TP-sharded inference column-parallel. + f_a_proj=TELinear, + f_b_proj=InferenceColumnParallelLinear, + g_a_proj=TELinear, + g_b_proj=InferenceColumnParallelLinear, + out_norm=TENorm, + out_proj=InferenceRowParallelLinear, + ), + ), + self_attn_bda=get_bias_dropout_add, + ), + ), # Started with spec from gpt_layer_specs.py (with MLP removed) # Using the TE spec because we had problems getting the non-TE spec # working @@ -413,6 +463,7 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): InferenceLayerNormColumnParallelLinear, InferenceRowParallelLinear ), gdn_layer=hybrid_inference_stack_spec.submodules.gdn_layer, + kda_layer=hybrid_inference_stack_spec.submodules.kda_layer, attention_layer=hybrid_inference_stack_spec.submodules.attention_layer, dsa_layer=hybrid_inference_stack_spec.submodules.dsa_layer, mlp_layer=hybrid_inference_stack_spec.submodules.mlp_layer, diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index 9ced52843d5..cfde29d7a4e 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -2,7 +2,7 @@ import logging from contextlib import nullcontext -from typing import Literal, Optional +from typing import Any, Callable, Literal, Optional import torch from torch import Tensor @@ -451,6 +451,8 @@ def forward( padding_mask: Optional[Tensor] = None, compute_mtp_loss: bool = True, cp_batch: ContextParallelBatch | None = None, + output_processor: Optional[Callable[..., Any]] = None, + output_processor_context: Optional[Any] = None, ) -> Tensor: """Forward function of the Hybrid model. This function passes the input tensors through the embedding layer, and then the decoder and finally into the post @@ -664,6 +666,26 @@ def forward( ), main_hidden_states=hidden_states, ) + # Match GPTModel's hook for caller-owned output projection and loss. + if output_processor is not None: + return output_processor( + hidden_states=hidden_states, + output_layer=self.output_layer, + output_weight=output_weight, + labels=labels, + loss_mask=loss_mask, + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + decoder_input=decoder_input, + inference_context=inference_context, + packed_seq_params=packed_seq_params, + runtime_gather_output=runtime_gather_output, + context=output_processor_context, + compute_language_model_loss=self.compute_language_model_loss, + scale_logits=self._scale_logits, + config=self.config, + ) sequence_parallel_override = False if ( in_inference_mode diff --git a/megatron/core/models/hybrid/layers/utils.py b/megatron/core/models/hybrid/layers/utils.py index fe5d99a6a62..45b4b077b56 100644 --- a/megatron/core/models/hybrid/layers/utils.py +++ b/megatron/core/models/hybrid/layers/utils.py @@ -1,6 +1,7 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. from megatron.core.ssm.gdn_layer_config import GDNLayerConfig +from megatron.core.ssm.kda_layer_config import KDALayerConfig from megatron.core.ssm.mamba_layer_config import MambaLayerConfig from megatron.core.ssm.mlp_layer_config import MLPLayerConfig from megatron.core.transformer.attention_layer_config import AttentionLayerConfig @@ -15,6 +16,7 @@ class Symbols: MAMBA = "M" GDN = 'G' + KDA = 'K' ATTENTION = "*" DS_ATTENTION = "D" MLA = "+" @@ -25,6 +27,7 @@ class Symbols: LAYER_CONFIG_MAP = { MAMBA: MambaLayerConfig, GDN: GDNLayerConfig, + KDA: KDALayerConfig, ATTENTION: AttentionLayerConfig, DS_ATTENTION: DSALayerConfig, MLA: MLALayerConfig, diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index 52e7dbaa516..58243adb200 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -1,5 +1,6 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. from dataclasses import dataclass +from typing import Literal import torch import torch.distributed as dist @@ -27,6 +28,7 @@ class PackedSeqParams: tokens_per_sample: int = None pad_between_seqs: bool = None cp_scatter_cache: object = None + cp_partition_mode: Literal["zigzag", "contiguous"] = "zigzag" def __post_init__(self): """Pre-compute seq_idx for Mamba mixer CUDA graph compatibility. @@ -67,3 +69,17 @@ def __post_init__(self): .to(torch.int32) .unsqueeze(0) # Add a batch dimension ) + + +def resolve_cp_group( + static_cp_group: dist.ProcessGroup, packed_seq_params: PackedSeqParams = None +) -> dist.ProcessGroup: + """Return the dynamic CP group from packed_seq_params when available, else the static one. + + Dynamic CP assigns a per-microbatch CP group that may differ from the + process-group stored at model construction time. This helper centralises + the resolution logic used by GPTModel, GatedDeltaNet, and MTP layers. + """ + if packed_seq_params is not None and packed_seq_params.cp_group is not None: + return packed_seq_params.cp_group + return static_cp_group diff --git a/megatron/core/ssm/gated_delta_net/__init__.py b/megatron/core/ssm/gated_delta_net/__init__.py index e5a5b4e6a9f..1046cd1ea8b 100644 --- a/megatron/core/ssm/gated_delta_net/__init__.py +++ b/megatron/core/ssm/gated_delta_net/__init__.py @@ -23,13 +23,21 @@ chunk_gdn2, torch_chunk_gdn2, ) +from megatron.core.ssm.gated_delta_net.kda import ( + HAVE_FLA_KDA, + KimiDeltaAttention, + KimiDeltaAttentionSubmodules, +) __all__ = [ "HAVE_FLA", "HAVE_FLA_GDN2", + "HAVE_FLA_KDA", "GatedDeltaNet", "GatedDeltaNet2", "GatedDeltaNetSubmodules", + "KimiDeltaAttention", + "KimiDeltaAttentionSubmodules", "causal_conv1d", "chunk_gated_delta_rule", "chunk_gdn2", diff --git a/megatron/core/ssm/gated_delta_net/common.py b/megatron/core/ssm/gated_delta_net/common.py index ea23b40111b..37840584d4e 100644 --- a/megatron/core/ssm/gated_delta_net/common.py +++ b/megatron/core/ssm/gated_delta_net/common.py @@ -10,11 +10,10 @@ import logging from dataclasses import dataclass from functools import lru_cache -from typing import Callable, Optional, Protocol, Union +from typing import Optional, Protocol, Union import torch import torch.nn as nn -import torch.nn.functional as F from megatron.core.fp8_utils import get_fp8_align_size from megatron.core.inference.contexts import BaseInferenceContext @@ -43,6 +42,7 @@ try: from fla.modules.convolution import causal_conv1d from fla.modules.l2norm import l2norm + from fla.ops.cp import build_cp_context from fla.ops.gated_delta_rule import chunk_gated_delta_rule HAVE_FLA = True @@ -50,6 +50,7 @@ causal_conv1d = None l2norm = None chunk_gated_delta_rule = None + build_cp_context = None HAVE_FLA = False @@ -146,8 +147,8 @@ def __init__( ignored; GDN implements context parallelism with its own all-to-alls rather than the attention CP communication schemes. pp_layer_offset: Offset of this pipeline stage's first global layer. + is_mtp_layer (bool): Whether this module is inside an MTP prediction depth. """ - del is_mtp_layer if not HAVE_FLA: raise ImportError( "FLA is not installed. Please install it with " @@ -159,6 +160,7 @@ def __init__( # Attributes from arguments self.layer_number = layer_number self.pp_layer_offset = pp_layer_offset + self.is_mtp_layer = is_mtp_layer self.bias = bias self.conv_bias = conv_bias self.conv_init = conv_init @@ -196,14 +198,16 @@ def __init__( "in_proj_extra_dim", "in_proj_split_names", "in_proj_split_sections", - "feat_dim_split", "gated_delta_rule", ) self._setup_variant_attrs() for attr in attrs_to_check: assert getattr(self, attr, None) is not None, f"Attribute {attr} for GDN is not set" - # QK, V, gate, shared across all variants - self.in_proj_qkvg_dim = self.qk_dim * 2 + self.v_dim * 2 + # Two-stage gates use separate projections; in_proj emits QKV only. + if getattr(self, "two_stage_gates", False): + self.in_proj_qkvg_dim = self.qk_dim * 2 + self.v_dim + else: + self.in_proj_qkvg_dim = self.qk_dim * 2 + self.v_dim * 2 self.in_proj_dim = self.in_proj_qkvg_dim + self.in_proj_extra_dim if self.config.fp8: @@ -249,7 +253,9 @@ def __init__( self.dt_bias = nn.Parameter( torch.empty( - self.dt_bias_dim, dtype=self.config.params_dtype, device=torch.cuda.current_device() + self.dt_bias_dim, + dtype=getattr(self, "gate_params_dtype", self.config.params_dtype), + device=torch.cuda.current_device(), ) ) setattr(self.dt_bias, "tensor_model_parallel", True) @@ -257,7 +263,9 @@ def __init__( self.A_log = nn.Parameter( torch.empty( - self.a_log_dim, dtype=self.config.params_dtype, device=torch.cuda.current_device() + self.a_log_dim, + dtype=getattr(self, "gate_params_dtype", self.config.params_dtype), + device=torch.cuda.current_device(), ) ) setattr(self.A_log, "tensor_model_parallel", True) @@ -272,8 +280,10 @@ def __init__( ) self.recompute_norm_out = False self.norm_out_checkpoint = None - if self.config.recompute_granularity == "selective": + self.recompute_gdn = False + if self.config.recompute_granularity == "selective" and self.config.recompute_modules: self.recompute_norm_out = "gdn_norm_out" in self.config.recompute_modules + self.recompute_gdn = "gdn" in self.config.recompute_modules self.out_proj = build_module( submodules.out_proj, @@ -289,6 +299,9 @@ def __init__( tp_group=self.pg_collection.tp, name=(name + ".out_proj") if name is not None else None, ) + # TODO: Packed sequence cu_seqlens can vary per batch; cache only static SBHD + # cp_context entries here and revisit routing metadata lifetime in the CP layout refactor. + self._chunkwise_cp_context_cache: dict[tuple[int, int], tuple[torch.Tensor, object]] = {} self.reset_parameters() @@ -393,6 +406,7 @@ def _prepare_input_for_gated_delta_rule( batch: int, seq_len: int, *gate_feats: tuple[torch.Tensor], + cp_size_headwise: int | None = None, ) -> dict[str, torch.Tensor]: """ Prepare all gated delta rule kernel inputs. @@ -406,11 +420,10 @@ def _prepare_input_for_gated_delta_rule( ``k``, ``v``, ``g``, plus the variant-specific gates), and the output gate (z) tensor under the ``gate`` key, which is not a kernel input. """ + cp_size = self.cp_size if cp_size_headwise is None else cp_size_headwise # Split qkv into query_key and value query_key, value = torch.split( - qkv, - [2 * self.qk_dim_local_tp // self.cp_size, self.v_dim_local_tp // self.cp_size], - dim=-1, + qkv, [2 * self.qk_dim_local_tp // cp_size, self.v_dim_local_tp // cp_size], dim=-1 ) # Reshape query_key and value @@ -422,7 +435,7 @@ def _prepare_input_for_gated_delta_rule( query_key = l2norm(query_key.contiguous()) # Split query and key - split_size = self.qk_dim_local_tp // self.key_head_dim // self.cp_size + split_size = self.qk_dim_local_tp // self.key_head_dim // cp_size query, key = torch.split(query_key, [split_size, split_size], dim=2) # Expand query and key if needed (grouped query attention) diff --git a/megatron/core/ssm/gated_delta_net/kda.py b/megatron/core/ssm/gated_delta_net/kda.py new file mode 100644 index 00000000000..b7fe670cf38 --- /dev/null +++ b/megatron/core/ssm/gated_delta_net/kda.py @@ -0,0 +1,813 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Kimi Delta Attention, a channel-wise Gated DeltaNet variant.""" + +import math +from dataclasses import dataclass +from functools import partial +from typing import Optional + +import torch +import torch.nn.functional as F + +from megatron.core import tensor_parallel +from megatron.core.context_parallel_layout import convert_module_input_tensors_cp_partition_mode +from megatron.core.fp8_utils import get_fp8_disabled_context +from megatron.core.inference.contexts import BaseInferenceContext +from megatron.core.jit import jit_fuser +from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.ssm.gated_delta_net.common import ( + HAVE_FLA, + GatedDeltaNetSubmodules, + _GDNBase, + a2a_cp_to_hp, + a2a_hp_to_cp, + build_cp_context, + causal_conv1d, + get_parameter_local_cp, +) +from megatron.core.transformer.identity_op import IdentityOp +from megatron.core.transformer.module import mark_keep_in_fp32 +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import deprecate_inference_params, nvtx_range_pop, nvtx_range_push + +try: + from fla.modules.fused_norm_gate import rms_norm_gated + from fla.ops.kda import chunk_kda + + # KDA also relies on the shared FLA convolution, normalization, and CP helpers. + HAVE_FLA_KDA = HAVE_FLA +except ImportError: # pragma: no cover + chunk_kda = None + HAVE_FLA_KDA = False + + +@dataclass +class KimiDeltaAttentionSubmodules(GatedDeltaNetSubmodules): + """Submodules used by Kimi Delta Attention.""" + + beta_proj: ModuleSpec | type = IdentityOp + + # Optional low-rank decay and output gates. + f_a_proj: ModuleSpec | type = IdentityOp + f_b_proj: ModuleSpec | type = IdentityOp + g_a_proj: ModuleSpec | type = IdentityOp + g_b_proj: ModuleSpec | type = IdentityOp + + +class KimiDeltaAttention(_GDNBase): + """Channel-wise Gated DeltaNet variant with direct Q/K/V/F/G projections. + + Supports direct or two-stage gates with equal query/key and value heads. + Recurrent inference is not implemented. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: KimiDeltaAttentionSubmodules, + layer_number: int | None = None, + bias: bool = False, + conv_bias: bool = False, + conv_init: float | None = None, + use_qk_l2norm: bool = True, + A_init_range: tuple[float, float] = (1, 16), + pg_collection: ProcessGroupCollection | None = None, + *, + name: str | None = None, + cp_comm_type: str | None = None, + pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, + ) -> None: + if not HAVE_FLA or not HAVE_FLA_KDA: # pragma: no cover + raise ImportError( + "FLA KDA is not installed. Install flash-linear-attention with KDA support." + ) + + self.two_stage_gates = config.kda_two_stage_gates + + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + bias=bias, + conv_bias=conv_bias, + conv_init=conv_init, + use_qk_l2norm=use_qk_l2norm, + A_init_range=A_init_range, + pg_collection=pg_collection, + name=name, + cp_comm_type=cp_comm_type, + pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, + ) + + # KDA keeps beta in a separate projection so its checkpoint layout remains + # independent from the direct Q/K/V/F/G projection. + self.beta_proj = build_module( + submodules.beta_proj, + self.hidden_size, + self.num_key_heads, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=bias, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name="beta_proj", + tp_group=self.pg_collection.tp, + name=(name + ".beta_proj") if name is not None else None, + ) + + # These kernel parameters participate in FP32 gate math and must remain + # FP32 through model casting, optimizer construction, and checkpointing. + mark_keep_in_fp32(self.dt_bias) + mark_keep_in_fp32(self.A_log) + + # Preserve the dev direct-gate path; low-rank gates are precomputed in FP32. + self.use_gate_in_kernel = not self.two_stage_gates + if self.two_stage_gates: + for prefix, head_dim, output_dim in ( + ("f", self.key_head_dim, self.qk_dim), + ("g", self.value_head_dim, self.v_dim), + ): + a_name, b_name = f"{prefix}_a_proj", f"{prefix}_b_proj" + setattr( + self, + a_name, + build_module( + getattr(submodules, a_name), + self.hidden_size, + head_dim, + config=self.config, + init_method=self.config.init_method, + bias=bias, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + name=f"{name}.{a_name}" if name is not None else None, + ), + ) + setattr( + self, + b_name, + build_module( + getattr(submodules, b_name), + head_dim, + output_dim, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=bias, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name=b_name, + tp_group=self.pg_collection.tp, + name=f"{name}.{b_name}" if name is not None else None, + ), + ) + + def _setup_variant_attrs(self) -> None: + """Set KDA dimensions, projection checkpoint metadata, and kernel callable.""" + + self.gdn_pre_gated_delta_rule_fusion = self.config.gdn_pre_gated_delta_rule_fusion + + # Channel-wise raw memory-decay gate g. + self.in_proj_extra_dim = self.qk_dim + + # Per-section sizes (and names) of the in_proj output, local to this TP rank. + # Used for the CP head permutation, post-a2a split, and sharded checkpoint split. + self.in_proj_split_names = ["query", "key", "value", "g", "gate"] + self.in_proj_split_sections = ( + self.qk_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + ) + if self.two_stage_gates: + self.in_proj_extra_dim = 0 + self.in_proj_split_names = self.in_proj_split_names[:3] + self.in_proj_split_sections = self.in_proj_split_sections[:3] + self.dt_bias_dim = self.qk_dim_local_tp + self.a_log_dim = self.num_k_heads_local_tp + self.gate_params_dtype = torch.float32 + self.gated_delta_rule = chunk_kda + + def _get_feat_dim_split(self, cp_size_headwise: int) -> tuple[int, int, int]: + """Return KDA qkv/raw-g/output-gate split sizes for runtime headwise CP.""" + + return ( + (self.qk_dim_local_tp * 2 + self.v_dim_local_tp) // cp_size_headwise, + self.qk_dim_local_tp // cp_size_headwise, + self.v_dim_local_tp // cp_size_headwise, + ) + + def _reset_dt_bias(self) -> None: + """Initialize the KDA channel-wise step-size bias in inverse-softplus space.""" + + dt_min, dt_max, dt_init_floor = 0.001, 0.1, 1e-4 + dt = torch.exp( + torch.rand(self.dt_bias_dim, device=self.dt_bias.device, dtype=self.dt_bias.dtype) + * (math.log(dt_max) - math.log(dt_min)) + + math.log(dt_min) + ).clamp(min=dt_init_floor) + self.dt_bias.data.copy_(dt + torch.log(-torch.expm1(-dt))) + + @jit_fuser + def _compute_gates( + self, + A_log_local_cp: torch.Tensor, + dt_bias_local_cp: torch.Tensor, + batch: int, + seq_len: int, + *gate_feats: tuple[torch.Tensor], + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Shape the raw channel-wise decay gate and activate the write strength.""" + + # ``gate_feats`` follows the KDA pre-GDR order: raw g, then separate beta. + raw_g, beta = gate_feats + num_key_heads = A_log_local_cp.numel() + raw_g = raw_g.reshape(batch, seq_len, num_key_heads, self.key_head_dim) + beta = beta.reshape(batch, seq_len, num_key_heads).float().sigmoid() + return raw_g, {"beta": beta.contiguous()} + + @jit_fuser + def _apply_gated_norm(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + """Apply per-head RMSNorm followed by KDA's sigmoid output gate.""" + + x_dtype = x.dtype + x = x.reshape(-1, self.value_head_dim) + gate = gate.reshape(-1, self.value_head_dim) + if self.two_stage_gates: + # Round only after both normalization and sigmoid gating. + weight = self.out_norm.weight + if self.config.layernorm_zero_centered_gamma: + weight = weight + 1 + return rms_norm_gated( + x, gate, weight, None, activation="sigmoid", eps=self.config.layernorm_epsilon + ) + x = self.out_norm(x) + return (x * torch.sigmoid(gate.float())).to(x_dtype) + + def _proj(self, proj, x): + """Apply a projection GEMM, forcing BF16 when kda_disable_fp8 is set. + + vLLM rollout keeps every KDA projection in BF16 (the checkpoint stores them + BF16 with no FP8 scales), so the actor must match to keep actor<->rollout + pearson sharp under FP8 hybrid training.""" + if self.config.kda_disable_fp8: + with get_fp8_disabled_context(self.config): + return proj(x) + return proj(x) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + inference_context: Optional[BaseInferenceContext] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + sequence_len_offset: Optional[int] = None, + *, + pg_collection: Optional[ProcessGroupCollection] = None, + inference_params: Optional[BaseInferenceContext] = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Run the direct-projection KDA training path.""" + + del attention_mask, sequence_len_offset, kwargs + inference_context = deprecate_inference_params(inference_context, inference_params) + + active_pg_collection = pg_collection if pg_collection is not None else self.pg_collection + base_cp_group = active_pg_collection.cp + cp_group = resolve_cp_group(base_cp_group, packed_seq_params) + if self.config.linear_cp_mode == "chunkwise": + cp_group_chunkwise = cp_group + cp_group_headwise = None + elif self.config.linear_cp_mode == "headwise": + cp_group_chunkwise = None + cp_group_headwise = cp_group + elif cp_group.size() == 1: + cp_group_chunkwise = None + cp_group_headwise = None + else: + raise ValueError( + f"Unsupported linear_cp_mode {self.config.linear_cp_mode!r}; " + "expected 'headwise' or 'chunkwise'." + ) + + cp_size_chunkwise = cp_group_chunkwise.size() if cp_group_chunkwise is not None else 1 + cp_size_headwise = cp_group_headwise.size() if cp_group_headwise is not None else 1 + cp_size_runtime = cp_group.size() + back_to_input_converter = None + if self.config.linear_cp_mode == "chunkwise": + hidden_states, back_to_input_converter = convert_module_input_tensors_cp_partition_mode( + hidden_states=hidden_states, + packed_seq_params=packed_seq_params, + cp_group=cp_group_chunkwise, + tp_group=self.tp_group, + tp_cp_group=getattr(active_pg_collection, "tp_cp", None), + target_partition_mode="contiguous", + sequence_parallel=self.config.sequence_parallel, + config=self.config, + ) + + seq_len_local, batch, _ = hidden_states.shape + seq_len_post_headwise = seq_len_local * self.sp_size * cp_size_headwise + seq_len_global = seq_len_post_headwise * cp_size_chunkwise + + if inference_context is not None: + assert ( + inference_context.is_static_batching() + ), "KimiDeltaAttention does not currently support dynamic inference batching." + assert not self.config.sequence_parallel + raise NotImplementedError("KimiDeltaAttention does not support inference for now.") + + if cp_size_headwise > 1 and ( + ( + packed_seq_params is not None + and packed_seq_params.qkv_format == "thd" + and packed_seq_params.cp_partition_mode != "zigzag" + ) + or ( + (packed_seq_params is None or packed_seq_params.qkv_format != "thd") + and self.config.cp_partition_mode != "zigzag" + ) + ): + raise ValueError( + "KimiDeltaAttention with headwise CP requires zigzag layout. CP partition " + "conversion must be handled before calling KimiDeltaAttention." + ) + + if packed_seq_params is not None and packed_seq_params.qkv_format == "thd": + if batch != 1: + raise ValueError("Packed KDA expects batch dimension to be 1.") + cu_seqlens_q = self._resolve_cu_seqlens( + packed_seq_params.cu_seqlens_q_padded, + packed_seq_params.cu_seqlens_q, + seq_len_global, + "cu_seqlens_q", + cp_size=cp_size_runtime, + ) + cu_seqlens_kv = self._resolve_cu_seqlens( + packed_seq_params.cu_seqlens_kv_padded, + packed_seq_params.cu_seqlens_kv, + seq_len_global, + "cu_seqlens_kv", + cp_size=cp_size_runtime, + ) + self._validate_packed_cu_seqlens(cu_seqlens_q, cu_seqlens_kv) + else: + cu_seqlens_q = None + + if cp_size_chunkwise > 1: + if cu_seqlens_q is None: + cache_key = (seq_len_global, batch) + cached = self._chunkwise_cp_context_cache.get(cache_key) + if cached is None: + cached_cu_seqlens = ( + torch.arange( + batch + 1, device=torch.cuda.current_device(), dtype=torch.long + ) + * seq_len_global + ) + cached_ctx = build_cp_context( + cu_seqlens=cached_cu_seqlens, + group=cp_group_chunkwise, + conv1d_kernel_size=self.conv_kernel_dim, + ) + cached = (cached_cu_seqlens, cached_ctx) + self._chunkwise_cp_context_cache[cache_key] = cached + cu_seqlens_q, chunkwise_cp_context = cached + else: + chunkwise_cp_context = build_cp_context( + cu_seqlens=cu_seqlens_q, + group=cp_group_chunkwise, + conv1d_kernel_size=self.conv_kernel_dim, + ) + else: + chunkwise_cp_context = None + + if self.recompute_gdn and self.training: + + def _checkpointed_compute(hidden_states): + return self._forward_compute( + hidden_states, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + packed_seq_params, + chunkwise_cp_context, + ) + + out, out_bias = tensor_parallel.checkpoint(_checkpointed_compute, False, hidden_states) + else: + out, out_bias = self._forward_compute( + hidden_states, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + packed_seq_params, + chunkwise_cp_context, + ) + + if back_to_input_converter is not None: + out = back_to_input_converter.convert( + out, seq_dim=0, sequence_parallel=self.config.sequence_parallel + ) + + return out, out_bias + + def _forward_compute( + self, + hidden_states, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + packed_seq_params, + chunkwise_cp_context, + ): + """Core KDA computation (in_proj -> conv1d -> gated_delta_rule -> norm -> out_proj).""" + + # Input projections. Beta intentionally remains a separate matrix. + nvtx_range_push(suffix="in_proj") + qkvfg, _ = self._proj(self.in_proj, hidden_states) + beta, _ = self._proj(self.beta_proj, hidden_states) + nvtx_range_pop(suffix="in_proj") + + qkvfg, thd_cp_a2a_inv = a2a_cp_to_hp( + qkvfg, + self.in_proj_split_sections, + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q, + seq_len_post_headwise, + packed_seq_params, + ) + beta, _ = a2a_cp_to_hp( + beta, + (self.num_k_heads_local_tp,), + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q, + seq_len_post_headwise, + packed_seq_params, + ) + + raw_g = gate = None + if self.two_stage_gates: + f_a, _ = self._proj(self.f_a_proj, hidden_states) + raw_g, _ = self._proj(self.f_b_proj, f_a) + g_a, _ = self._proj(self.g_a_proj, hidden_states) + gate, _ = self._proj(self.g_b_proj, g_a) + raw_g, _ = a2a_cp_to_hp( + raw_g, + (self.qk_dim_local_tp,), + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q, + seq_len_post_headwise, + packed_seq_params, + ) + gate, _ = a2a_cp_to_hp( + gate, + (self.v_dim_local_tp,), + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q, + seq_len_post_headwise, + packed_seq_params, + ) + + if self.gdn_pre_gated_delta_rule_fusion: + raise NotImplementedError( + "gdn_pre_gated_delta_rule_fusion is not implemented for KDA yet." + ) + + if cp_size_chunkwise > 1 and packed_seq_params is None and batch > 1: + raise ValueError( + "KDA chunkwise CP with SBHD inputs currently requires micro_batch_size == 1 " + "when cp_context is used. Use packed THD input or micro_batch_size=1." + ) + if cp_size_chunkwise > 1 and self.config.gdn_conv_pad_alignment is not None: + raise ValueError( + "gdn_conv_pad_alignment is incompatible with KDA chunkwise CP. Padding " + "chunk-local causal-conv inputs can change later chunk numerics." + ) + + nvtx_range_push(suffix="pre_gated_delta_rule") + query, key, value, gate, beta, raw_g, A_log, dt_bias = self.pre_gated_delta_rule( + qkvfg, + beta, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q, + chunkwise_cp_context, + packed_seq_params=packed_seq_params, + raw_g=raw_g, + gate=gate, + ) + kernel_inputs = {"q": query, "k": key, "v": value, "g": raw_g, "beta": beta} + if self.use_gate_in_kernel: + kernel_inputs["A_log"] = A_log + kernel_inputs["dt_bias"] = dt_bias + nvtx_range_pop(suffix="pre_gated_delta_rule") + + nvtx_range_push(suffix="gated_delta_rule") + core_attn_out, _ = self.gated_delta_rule( + **kernel_inputs, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=self.two_stage_gates and self.use_qk_l2norm, + use_gate_in_kernel=self.use_gate_in_kernel, + safe_gate=self.config.kda_safe_gate, + lower_bound=self.config.kda_lower_bound, + state_v_first=self.two_stage_gates, + cu_seqlens=cu_seqlens_q, + cp_context=chunkwise_cp_context, + ) + nvtx_range_pop(suffix="gated_delta_rule") + + if self.recompute_norm_out and self.training: + self.norm_out_checkpoint = tensor_parallel.CheckpointWithoutOutput() + norm_func = partial( + self._gated_norm_and_layout_restore, + thd_cp_a2a_inv=thd_cp_a2a_inv, + batch=batch, + seq_len=seq_len_post_headwise, + packed_seq_params=packed_seq_params, + cp_size_headwise=cp_size_headwise, + cp_group_headwise=cp_group_headwise, + cp_size_chunkwise=cp_size_chunkwise, + cp_group_chunkwise=cp_group_chunkwise, + cu_seqlens_q=cu_seqlens_q, + ) + norm_out = self.norm_out_checkpoint.checkpoint(norm_func, core_attn_out, gate) + else: + norm_out = self._gated_norm_and_layout_restore( + core_attn_out, + gate, + thd_cp_a2a_inv, + batch, + seq_len_post_headwise, + packed_seq_params, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + ) + + # Output projection. + nvtx_range_push(suffix="out_proj") + out, out_bias = self._proj(self.out_proj, norm_out) + nvtx_range_pop(suffix="out_proj") + + if self.recompute_norm_out and self.training: + self.norm_out_checkpoint.discard_output_and_register_recompute(out) + + return out, out_bias + + def _gated_norm_and_layout_restore( + self, + core_attn_out: torch.Tensor, + gate: torch.Tensor, + thd_cp_a2a_inv: torch.Tensor | None, + batch: int, + seq_len: int, + packed_seq_params: PackedSeqParams | None, + cp_size_headwise: int, + cp_group_headwise: torch.distributed.ProcessGroup | None, + cp_size_chunkwise: int, + cp_group_chunkwise: torch.distributed.ProcessGroup | None, + cu_seqlens_q: torch.Tensor | None, + ) -> torch.Tensor: + """Apply KDA's gated output norm and restore its context-parallel layout.""" + + nvtx_range_push(suffix="gated_norm") + norm_out = self._apply_gated_norm(core_attn_out, gate) + nvtx_range_pop(suffix="gated_norm") + + norm_out = norm_out.reshape(batch, seq_len, -1) + norm_out = norm_out.transpose(0, 1).contiguous() + + return a2a_hp_to_cp( + norm_out, cp_size_headwise, cp_group_headwise, packed_seq_params, thd_cp_a2a_inv + ) + + def pre_gated_delta_rule( + self, + qkvfg, + beta, + batch, + seq_len, + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q=None, + chunkwise_cp_context=None, + packed_seq_params=None, + *, + raw_g=None, + gate=None, + ): + """Prepare QKV, output gate, beta, and raw decay tensors before KDA.""" + + qkvfg = qkvfg.transpose(0, 1) + beta = beta.transpose(0, 1) + if self.two_stage_gates: + qkv = qkvfg + raw_g = raw_g.transpose(0, 1) + gate = gate.transpose(0, 1) + else: + qkv, raw_g, gate = torch.split( + qkvfg, self._get_feat_dim_split(cp_size_headwise), dim=-1 + ) + gate = gate.reshape(batch, seq_len, -1, self.value_head_dim) + + nvtx_range_push(suffix="conv1d") + kernel_seq_len = qkv.shape[1] + qkv_channels_split_sections = [ + self.qk_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + ] + conv1d_weight = get_parameter_local_cp( + self.conv1d.weight, + dim=0, + cp_group=cp_group_headwise, + split_sections=qkv_channels_split_sections, + ) + conv1d_bias = ( + get_parameter_local_cp( + self.conv1d.bias, + dim=0, + cp_group=cp_group_headwise, + split_sections=qkv_channels_split_sections, + ) + if self.conv_bias + else None + ) + if self.config.deterministic_mode or causal_conv1d is None: + qkv = qkv.transpose(1, 2).contiguous() + conv_out = F.conv1d( + input=qkv, + weight=conv1d_weight, + bias=conv1d_bias, + stride=self.conv1d.stride, + padding=self.conv1d.padding, + dilation=self.conv1d.dilation, + groups=self.conv_dim_local_tp // cp_size_headwise, + ) + qkv = self.act_fn(conv_out[..., :kernel_seq_len]) + qkv = qkv.transpose(1, 2).contiguous() + else: + if self.activation not in ("silu", "swish"): + raise ValueError(f"FLA causal convolution requires SiLU, got {self.activation}.") + orig_seq = qkv.shape[1] + pad_n = 0 + conv_input = qkv.contiguous() + conv_cu_seqlens = cu_seqlens_q + conv_cp_context = chunkwise_cp_context + if self.config.gdn_conv_pad_alignment is not None: + if packed_seq_params is None or cu_seqlens_q is None: + raise ValueError( + "gdn_conv_pad_alignment is only supported with packed sequence " + "parameters in THD format. SBHD inputs do not need causal-conv padding." + ) + if chunkwise_cp_context is not None: + raise ValueError( + "gdn_conv_pad_alignment is incompatible with KDA chunkwise CP. Padding " + "chunk-local causal-conv inputs can change later chunk numerics." + ) + pad_n = -orig_seq % self.config.gdn_conv_pad_alignment + if pad_n > 0: + conv_input = torch.nn.functional.pad(conv_input, (0, 0, 0, pad_n)) + conv_cu_seqlens = cu_seqlens_q.clone() + conv_cu_seqlens[-1] += pad_n + qkv, _ = causal_conv1d( + x=conv_input, + weight=conv1d_weight.squeeze(1), + bias=conv1d_bias, + activation=self.activation, + initial_state=None, + output_final_state=False, + cu_seqlens=conv_cu_seqlens, + cp_context=conv_cp_context, + ) + if pad_n > 0: + qkv = qkv[:, :orig_seq, :] + nvtx_range_pop(suffix="conv1d") + + A_log_local_cp = get_parameter_local_cp(self.A_log, dim=0, cp_group=cp_group_headwise) + dt_bias_local_cp = get_parameter_local_cp(self.dt_bias, dim=0, cp_group=cp_group_headwise) + + if self.two_stage_gates: + query_key, value = torch.split( + qkv, + [ + 2 * self.qk_dim_local_tp // cp_size_headwise, + self.v_dim_local_tp // cp_size_headwise, + ], + dim=-1, + ) + query_key = query_key.reshape(batch, seq_len, -1, self.key_head_dim) + value = value.reshape(batch, seq_len, -1, self.value_head_dim) + query, key = query_key.chunk(2, dim=2) + repeat_factor = self.num_value_heads // self.num_key_heads + if repeat_factor > 1: + query = query.repeat_interleave(repeat_factor, dim=2) + key = key.repeat_interleave(repeat_factor, dim=2) + num_key_heads = A_log_local_cp.numel() + raw_g = raw_g.reshape(batch, seq_len, num_key_heads, self.key_head_dim) + beta = beta.reshape(batch, seq_len, num_key_heads).float().sigmoid() + from fla.ops.kda.gate import fused_kda_gate + + raw_g = fused_kda_gate( + raw_g, + A_log_local_cp.contiguous(), + dt_bias=dt_bias_local_cp.contiguous(), + lower_bound=self.config.kda_lower_bound if self.config.kda_safe_gate else None, + ) + return ( + query.contiguous(), + key.contiguous(), + value.contiguous(), + gate, + beta.contiguous(), + raw_g.contiguous(), + A_log_local_cp, + dt_bias_local_cp, + ) + + nvtx_range_push(suffix="prepare_input_for_gated_delta_rule") + kernel_inputs = self._prepare_input_for_gated_delta_rule( + qkv, + gate, + A_log_local_cp, + dt_bias_local_cp, + batch, + kernel_seq_len, + raw_g, + beta, + cp_size_headwise=cp_size_headwise, + ) + nvtx_range_pop(suffix="prepare_input_for_gated_delta_rule") + + gate = kernel_inputs.pop("gate") + + return ( + kernel_inputs["q"], + kernel_inputs["k"], + kernel_inputs["v"], + gate, + kernel_inputs["beta"], + kernel_inputs["g"], + A_log_local_cp, + dt_bias_local_cp, + ) + + @staticmethod + def _validate_packed_cu_seqlens( + cu_seqlens_q: torch.Tensor, cu_seqlens_kv: torch.Tensor + ) -> None: + """Validate the self-attention boundary contract for packed KDA.""" + + if cu_seqlens_q.numel() < 2 or cu_seqlens_kv.numel() < 2: + raise ValueError( + "Packed KDA requires at least one sequence in both Q and KV boundaries." + ) + if cu_seqlens_q.shape != cu_seqlens_kv.shape or not torch.equal( + cu_seqlens_q, cu_seqlens_kv + ): + raise ValueError( + "Packed KDA requires cu_seqlens_q to equal cu_seqlens_kv, " + f"but got shapes {tuple(cu_seqlens_q.shape)} and " + f"{tuple(cu_seqlens_kv.shape)}." + ) + + def backward_dw(self) -> None: + """Execute weight-gradient computation for all KDA projections.""" + + super().backward_dw() + self.beta_proj.backward_dw() + if self.two_stage_gates: + self.f_a_proj.backward_dw() + self.f_b_proj.backward_dw() + self.g_a_proj.backward_dw() + self.g_b_proj.backward_dw() diff --git a/megatron/core/ssm/kda_layer_config.py b/megatron/core/ssm/kda_layer_config.py new file mode 100644 index 00000000000..9f3c1c74b27 --- /dev/null +++ b/megatron/core/ssm/kda_layer_config.py @@ -0,0 +1,10 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from megatron.core.transformer.transformer_config import TransformerConfig + + +class KDALayerConfig(TransformerConfig): + """Configuration for a KDA (Kimi Delta Attention) layer in a hybrid stack. + + Due to backwards-compatibility, this config's arguments are defined in TransformerConfig. + """ diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 0ead26dd15b..c8c4eae5539 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -414,6 +414,27 @@ class TransformerConfig(ModelParallelConfig): linear_num_value_heads: Optional[int] = 32 """Number of value and gate heads for the gated delta net.""" + kda_disable_fp8: bool = False + """Force KDA projections to BF16 even under FP8 training, + (KDA projections are BF16 in the checkpoint).""" + + kda_safe_gate: bool = False + """Whether the KDA kernel should use bounded gate values.""" + + kda_lower_bound: Optional[float] = None + """Optional lower bound for KDA's bounded gate values.""" + + kda_two_stage_gates: bool = False + """Use low-rank f_b(f_a(x)) and g_b(g_a(x)) gates with a QKV-only input projection.""" + + gdn_pre_gated_delta_rule_fusion: bool = False + """Whether to use the streamed Triton fusion for GatedDeltaNet pre-GDR preprocessing.""" + + gdn_conv_pad_alignment: Optional[int] = None + """When set, pad packed GDN causal-conv inputs to this token alignment. + This is only valid without chunkwise CP: padding a chunk-local causal-conv input changes + the sequence seen by later chunks and therefore changes the GDN recurrence numerics.""" + #################### # initialization #################### diff --git a/megatron/core/utils.py b/megatron/core/utils.py index caad870f120..31412cdb9fb 100644 --- a/megatron/core/utils.py +++ b/megatron/core/utils.py @@ -2814,6 +2814,16 @@ def nvtx_range_pop(msg=None, suffix=None) -> None: torch.cuda.nvtx.range_pop() +@contextmanager +def nvtx_range(msg=None, suffix=None): + """Create an NVTX range controlled by ``configure_nvtx_profiling``.""" + nvtx_range_push(msg, suffix) + try: + yield + finally: + nvtx_range_pop(msg, suffix) + + @lru_cache(maxsize=None) def _nvtx_decorator_get_func_path(func): """Get the path of a function. diff --git a/tests/unit_tests/ssm/test_kda_gate_precision.py b/tests/unit_tests/ssm/test_kda_gate_precision.py new file mode 100644 index 00000000000..856fa9ba6b5 --- /dev/null +++ b/tests/unit_tests/ssm/test_kda_gate_precision.py @@ -0,0 +1,115 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""KDA's FP32 gate parameters must survive BF16 model construction and casting.""" + +import pytest +import torch + +from megatron.core import parallel_state +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.ssm.gated_delta_net.kda import HAVE_FLA_KDA +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.module import Float16Module +from megatron.core.transformer.transformer_config import TransformerConfig +from tests.unit_tests.test_utilities import Utils + + +@pytest.mark.skipif(not HAVE_FLA_KDA, reason="FLA KDA is not installed.") +@pytest.mark.parametrize("variant", ["kda", "kda_direct", "gdn"]) +def test_gate_parameter_precision_through_bf16_wrapper(variant): + Utils.initialize_model_parallel(tensor_model_parallel_size=1, pipeline_model_parallel_size=1) + try: + model_parallel_cuda_manual_seed(123) + config = TransformerConfig( + num_layers=1, + hidden_size=256, + num_attention_heads=2, + linear_num_key_heads=2, + linear_num_value_heads=2, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_conv_kernel_dim=4, + params_dtype=torch.bfloat16, + bf16=True, + normalization="RMSNorm", + activation_func=torch.nn.functional.silu, + kda_two_stage_gates=variant == "kda", + kda_safe_gate=True, + kda_lower_bound=-5.0, + perform_initialization=True, + ) + layer_spec = getattr( + hybrid_stack_spec.submodules, f"{variant.removesuffix('_direct')}_layer" + ).submodules.self_attention + layer = layer_spec.module( + config=config, + submodules=layer_spec.submodules, + layer_number=1, + pg_collection=ProcessGroupCollection( + tp=parallel_state.get_tensor_model_parallel_group(), + cp=parallel_state.get_context_parallel_group(), + ), + ) + expected_dtype = torch.bfloat16 if variant == "gdn" else torch.float32 + expected = {} + for name in ("A_log", "dt_bias"): + param = getattr(layer, name) + assert param.dtype == expected_dtype + with torch.no_grad(): + param.fill_(0.12345678) + expected[name] = param.detach().clone() + + Float16Module(config, layer) + + assert layer.in_proj.weight.dtype == torch.bfloat16 + for name, reference in expected.items(): + param = getattr(layer, name) + assert param.dtype == expected_dtype + assert param.tensor_model_parallel is True + assert param.partition_dim == 0 + torch.testing.assert_close(param, reference, rtol=0, atol=0) + + if variant == "kda": + torch.manual_seed(42) + x = torch.randn(256, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + gate = torch.randn_like(x, requires_grad=True) + with torch.no_grad(): + layer.out_norm.weight.uniform_(0.5, 1.5) + actual = layer._apply_gated_norm(x, gate) + x_fp32 = x.float() + expected_norm = x_fp32 * torch.rsqrt( + x_fp32.square().mean(dim=-1, keepdim=True) + config.layernorm_epsilon + ) + expected_output = ( + expected_norm * layer.out_norm.weight.float() * gate.float().sigmoid() + ).to(x.dtype) + assert (actual.float() - expected_output.float()).abs().mean() < 2e-5 + actual.float().sum().backward() + for tensor in (x, gate, layer.out_norm.weight): + assert tensor.grad is not None + assert torch.isfinite(tensor.grad).all() + + if variant.startswith("kda"): + assert layer.use_gate_in_kernel == (variant == "kda_direct") + hidden = torch.randn( + 260, 1, 256, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + cu_seqlens = torch.tensor([0, 129, 260], device="cuda", dtype=torch.int32) + packed = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=131, + max_seqlen_kv=131, + ) + packed_output, _ = layer(hidden, attention_mask=None, packed_seq_params=packed) + separate_output = torch.cat( + [layer(segment, attention_mask=None)[0] for segment in hidden.split([129, 131])] + ) + torch.testing.assert_close(packed_output, separate_output, rtol=0.03, atol=0.002) + packed_output.float().square().mean().backward() + assert hidden.grad is not None and torch.isfinite(hidden.grad).all() + finally: + Utils.destroy_model_parallel() From 9bbe8cc4a1f7a0c9573d9522a92cf362b44812ea Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Thu, 10 Sep 2026 00:52:07 +0300 Subject: [PATCH 4/6] mHC Signed-off-by: Hollow Man --- megatron/core/models/hybrid/hybrid_block.py | 39 +++++---- .../hybrid/layers/hybrid_hyper_connection.py | 2 + megatron/core/recompute.py | 18 +++- megatron/core/transformer/hyper_connection.py | 37 +++++---- .../transformer/multi_token_prediction.py | 39 +++++---- .../core/transformer/transformer_config.py | 16 ++-- tests/unit_tests/models/test_hybrid_mhc.py | 64 +++++++++++++- .../test_hyper_connection_recompute.py | 11 --- .../transformer/test_mhc_precision.py | 83 +++++++++++++++++++ 9 files changed, 237 insertions(+), 72 deletions(-) create mode 100644 tests/unit_tests/transformer/test_mhc_precision.py diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py index 6a69452d345..7674ab354f3 100644 --- a/megatron/core/models/hybrid/hybrid_block.py +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -324,14 +324,15 @@ def __init__( if self.config.enable_mhc_connections and self.post_process and not self.is_mtp_layer: hc_mult = self.config.mhc_num_residual_streams hc_dim = self.config.hidden_size * hc_mult - self.hc_head_fn = mark_keep_in_fp32(nn.Parameter(torch.randn(hc_mult, hc_dim))) - self.hc_head_base = mark_keep_in_fp32(nn.Parameter(torch.zeros(hc_mult))) - self.hc_head_scale = mark_keep_in_fp32(nn.Parameter(torch.ones(1))) - nn.init.xavier_uniform_(self.hc_head_fn) - if self.config.sequence_parallel: - setattr(self.hc_head_fn, 'sequence_parallel', True) - setattr(self.hc_head_base, 'sequence_parallel', True) - setattr(self.hc_head_scale, 'sequence_parallel', True) + if self.config.mhc_learned_output_contract: + self.hc_head_fn = mark_keep_in_fp32(nn.Parameter(torch.randn(hc_mult, hc_dim))) + self.hc_head_base = mark_keep_in_fp32(nn.Parameter(torch.zeros(hc_mult))) + self.hc_head_scale = mark_keep_in_fp32(nn.Parameter(torch.ones(1))) + nn.init.xavier_uniform_(self.hc_head_fn) + if self.config.sequence_parallel: + setattr(self.hc_head_fn, 'sequence_parallel', True) + setattr(self.hc_head_base, 'sequence_parallel', True) + setattr(self.hc_head_scale, 'sequence_parallel', True) @property def layer_type_list(self) -> list[str]: @@ -647,14 +648,20 @@ def get_inner_quant_context(config, layer_number): if self.config.enable_mhc_connections and self.post_process and not self.is_mtp_layer: if (self.config.mtp_num_layers or 0) > 0: mhc_multistream = hidden_states - hidden_states = learned_output_contract( - hidden_states, - self.hc_head_fn, - self.hc_head_base, - self.hc_head_scale, - self.config.mhc_num_residual_streams, - self.config.layernorm_epsilon, - ) + if self.config.mhc_learned_output_contract: + hidden_states = learned_output_contract( + hidden_states, + self.hc_head_fn, + self.hc_head_base, + self.hc_head_scale, + self.config.mhc_num_residual_streams, + self.config.layernorm_epsilon, + ) + else: + n = self.config.mhc_num_residual_streams + hidden_states = hidden_states.unflatten(-1, (n, self.config.hidden_size)).mean( + dim=-2 + ) # Final layer norm. if self.post_process and self.post_layer_norm: diff --git a/megatron/core/models/hybrid/layers/hybrid_hyper_connection.py b/megatron/core/models/hybrid/layers/hybrid_hyper_connection.py index 55f4209b816..f6671dc8d82 100644 --- a/megatron/core/models/hybrid/layers/hybrid_hyper_connection.py +++ b/megatron/core/models/hybrid/layers/hybrid_hyper_connection.py @@ -29,6 +29,8 @@ class HyperConnectionHybridLayer(MegatronModule): switch between mHC-enabled and ordinary HybridStacks without key migration. """ + supports_hybrid_recompute_kwargs = True + def __init__(self, config: TransformerConfig, layer: MegatronModule) -> None: super().__init__(config=config) self.inner_layer = layer diff --git a/megatron/core/recompute.py b/megatron/core/recompute.py index 68974efb05b..572f37a5075 100644 --- a/megatron/core/recompute.py +++ b/megatron/core/recompute.py @@ -103,9 +103,9 @@ def custom_forward( else: inner_quantization_context = nullcontext() - # Build the full TransformerLayer kwarg set; for non-TL - # layers (currently MambaLayer in HybridStack) pop the kwargs - # they don't accept and treat the return as a single tensor. + # Build the full TransformerLayer kwarg set. Hybrid mHC wrappers expose + # an explicit capability flag so this module does not need to import + # hybrid_block (which would create a circular import). layer_kwargs = dict( hidden_states=hidden_states, attention_mask=attention_mask, @@ -120,6 +120,18 @@ def custom_forward( with inner_quantization_context: if isinstance(layer, TransformerLayer): hidden_states, context = layer(**layer_kwargs) + elif getattr(layer, "supports_hybrid_recompute_kwargs", False): + # HyperConnectionHybridLayer accepts the routing metadata + # consumed by wrapped MoE layers, but not cross-attention kwargs + # from the TransformerLayer interface. This also covers a wrapper + # around a MambaLayer; the wrapper narrows kwargs for its inner layer. + for k in ("context", "context_mask", "attention_bias"): + layer_kwargs.pop(k, None) + if packed_sequence_cp_metadata is not None: + layer_kwargs["packed_sequence_cp_metadata"] = ( + packed_sequence_cp_metadata + ) + hidden_states, context = layer(**layer_kwargs) else: # MambaLayer (HybridStack `M` slot) for k in ("context", "context_mask", "attention_bias", "padding_mask"): layer_kwargs.pop(k, None) diff --git a/megatron/core/transformer/hyper_connection.py b/megatron/core/transformer/hyper_connection.py index 6599a9308d6..ce0cd964550 100644 --- a/megatron/core/transformer/hyper_connection.py +++ b/megatron/core/transformer/hyper_connection.py @@ -102,7 +102,7 @@ def native_sinkhorn(input_logits: Tensor, num_iterations: int, eps: float = 1e-6 @torch.compile def native_h_aggregate(x: Tensor, h_pre: Tensor) -> Tensor: """Native n-stream weighted aggregation: out = sum_j(h_pre_j * x_j).""" - return (x * h_pre.unsqueeze(-1)).sum(dim=2) + return (x * h_pre.unsqueeze(-1)).sum(dim=2).to(x.dtype) @torch.compile @@ -112,19 +112,23 @@ def native_h_post_bda( """Native H_res.T @ residual + H_post * (x [+ bias]).""" s, b, n, C = original_residual.shape h_res_batched = h_res.view(s * b, n, n) - residual_batched = original_residual.view(s * b, n, C) + residual_batched = original_residual.view(s * b, n, C).to(h_res.dtype) mixed = torch.bmm(h_res_batched.transpose(1, 2), residual_batched).view(s, b, n, C) x_expanded = h_post.unsqueeze(-1) * x.unsqueeze(2) if bias is not None: bias_expanded = h_post.unsqueeze(-1) * bias.view(1, 1, 1, C) - return x_expanded + bias_expanded + mixed - return x_expanded + mixed + return (x_expanded + bias_expanded + mixed).to(original_residual.dtype) + return (x_expanded + mixed).to(original_residual.dtype) @torch.compile -def native_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6) -> Tuple[Tensor, Tensor]: +def native_proj_rms( + x: Tensor, weight: Tensor, eps: float = 1e-6, eps_inside_sqrt: bool = False +) -> Tuple[Tensor, Tensor]: """Native fused projection + RMS normalization.""" proj = torch.matmul(x, weight.t()) + if eps_inside_sqrt: + return proj, torch.rsqrt(x.square().mean(dim=-1, keepdim=True) + eps) norm = x.norm(dim=-1, keepdim=True) K = x.shape[-1] v = norm / math.sqrt(K) + eps @@ -239,7 +243,7 @@ def __init__(self, config: TransformerConfig, layer_number: int): mark_keep_in_fp32(self.alpha_post) mark_keep_in_fp32(self.alpha_res) mark_keep_in_fp32(self.bias) - self.norm_eps = 1e-6 + self.norm_eps = config.layernorm_epsilon if config.mhc_norm_eps_inside_sqrt else 1e-6 # Choose implementation: unified fused kernels vs reference modules. # The fused public API selects the backend per operation internally. @@ -251,9 +255,13 @@ def __init__(self, config: TransformerConfig, layer_number: int): # The fused path computes the projection and compute_h in one op, so # _projection_and_get_norm — and therefore _proj_rms_op — is only ever # reached on the unfused path. - self._proj_rms_op = native_proj_rms + self._proj_rms_op = partial( + native_proj_rms, eps_inside_sqrt=config.mhc_norm_eps_inside_sqrt + ) if config.use_fused_mhc: + if config.mhc_norm_eps_inside_sqrt or config.mhc_keep_mappings_in_fp32: + raise ValueError("Fused mHC does not support FP32 mixing or epsilon inside sqrt.") from megatron.core.fusions.fused_mhc_kernels import ( fused_h_aggregate, fused_h_post_bda, @@ -302,9 +310,7 @@ def _projection_and_get_norm(self, x: Tensor) -> Tuple[Tensor, Tensor]: x: [s, b, n*C] - n-stream hidden states """ s, b, nC = x.shape - # The mHC mapping computation runs in FP32: the parameters are kept in - # FP32 and the activations are upcast here, then compute_mappings casts - # the bounded mixing weights back to the activation dtype. + # Mapping projections use FP32 regardless of the activation dtype. x_2d = x.reshape(s * b, nC).to(torch.float32) weight = self.mapping_proj.weight.to(torch.float32) proj, r = self._proj_rms_op(x_2d, weight, self.norm_eps) @@ -389,10 +395,7 @@ def compute_mappings(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor]: h_res, self.sinkhorn_iterations, self.sinkhorn_eps ) # [s, b, n, n] - # The mixing weights are bounded (sigmoid outputs / doubly stochastic - # matrix), so after the FP32 computation they are safe to apply to the - # streams in the activation dtype. - dtype = x.dtype + dtype = torch.float32 if self.config.mhc_keep_mappings_in_fp32 else x.dtype return h_pre.to(dtype), h_post.to(dtype), h_res.to(dtype) @torch.compile @@ -512,7 +515,7 @@ def apply_h_res(self, h_res: Tensor, residual: Tensor) -> Tensor: # Reshape for bmm: [s, b, n, n] -> [s*b, n, n] h_res_batched = h_res.view(s * b, n, n) # [s, b, n*C] -> [s, b, n, C] -> [s*b, n, C] - residual_batched = residual.view(s, b, n, C).view(s * b, n, C) + residual_batched = residual.view(s * b, n, C).to(h_res.dtype) # Batch matrix multiply: [s*b, n, n].T @ [s*b, n, C] -> [s*b, n, C] mixed = torch.bmm(h_res_batched.transpose(1, 2), residual_batched) @@ -772,7 +775,7 @@ def _fused_h_res_h_post_bda_native( bda_func = get_bias_dropout_add(training, fused) with torch.cuda.nvtx.range("HyperConnection::bda"): output = bda_func((x_expanded, bias_expanded), mixed, dropout_prob) - return output + return output.to(original_residual.dtype) @nvtx_decorator(message="HyperConnection::fused_h_res_h_post_bda_with_checkpoint") def _fused_h_res_h_post_bda_with_checkpoint( @@ -847,7 +850,7 @@ def _native_wrapper(h_res, original_residual, h_post, x, *optional_bias): bias_expanded = None with torch.cuda.nvtx.range("HyperConnection::bda"): output = bda_func((x_expanded, bias_expanded), mixed, dropout_prob) - return output + return output.to(original_residual.dtype) ckpt = CheckpointWithoutOutput(ckpt_manager=manager) if has_bias: diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py index a844d8f76d8..925338baf37 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -1417,14 +1417,15 @@ def __init__( if self.mhc_enabled: hc_mult = self.config.mhc_num_residual_streams hc_dim = self.config.hidden_size * hc_mult - self.hc_head_fn = mark_keep_in_fp32(nn.Parameter(torch.randn(hc_mult, hc_dim))) - self.hc_head_base = mark_keep_in_fp32(nn.Parameter(torch.zeros(hc_mult))) - self.hc_head_scale = mark_keep_in_fp32(nn.Parameter(torch.ones(1))) - nn.init.xavier_uniform_(self.hc_head_fn) - if self.config.sequence_parallel: - setattr(self.hc_head_fn, "sequence_parallel", True) - setattr(self.hc_head_base, "sequence_parallel", True) - setattr(self.hc_head_scale, "sequence_parallel", True) + if self.config.mhc_learned_output_contract: + self.hc_head_fn = mark_keep_in_fp32(nn.Parameter(torch.randn(hc_mult, hc_dim))) + self.hc_head_base = mark_keep_in_fp32(nn.Parameter(torch.zeros(hc_mult))) + self.hc_head_scale = mark_keep_in_fp32(nn.Parameter(torch.ones(1))) + nn.init.xavier_uniform_(self.hc_head_fn) + if self.config.sequence_parallel: + setattr(self.hc_head_fn, "sequence_parallel", True) + setattr(self.hc_head_base, "sequence_parallel", True) + setattr(self.hc_head_scale, "sequence_parallel", True) self.offload_context = nullcontext() def get_inner_quantization_context(self) -> AbstractContextManager: @@ -1674,14 +1675,20 @@ def _postprocess(self, hidden_states: torch.Tensor): """ if self.mhc_enabled: - hidden_states = learned_output_contract( - hidden_states, - self.hc_head_fn, - self.hc_head_base, - self.hc_head_scale, - self.config.mhc_num_residual_streams, - self.config.layernorm_epsilon, - ) + if self.config.mhc_learned_output_contract: + hidden_states = learned_output_contract( + hidden_states, + self.hc_head_fn, + self.hc_head_base, + self.hc_head_scale, + self.config.mhc_num_residual_streams, + self.config.layernorm_epsilon, + ) + else: + n = self.config.mhc_num_residual_streams + hidden_states = hidden_states.unflatten(-1, (n, self.config.hidden_size)).mean( + dim=-2 + ) # Layer norm before shared head layer. hidden_states = apply_module(self.final_layernorm)(hidden_states) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index c8c4eae5539..2a3e86ed2b1 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1220,6 +1220,15 @@ class TransformerConfig(ModelParallelConfig): mhc_init_gating_factor: float = 0.01 """Initial value of Gating Factor (alpha in paper).""" + mhc_norm_eps_inside_sqrt: bool = False + """Use rsqrt(mean(x**2) + layernorm_epsilon) for the mHC mapping norm.""" + + mhc_keep_mappings_in_fp32: bool = False + """Keep mHC coefficients and stream mixing in FP32 until the output cast.""" + + mhc_learned_output_contract: bool = True + """Use learned hc_head_* weights to contract residual streams; otherwise take their mean.""" + use_fused_mhc: bool = False """Use fused kernels for mHC operations when supported. @@ -2289,13 +2298,6 @@ def __post_init__(self): if self.mhc_fused_backend != "auto" and not self.use_fused_mhc: raise ValueError("mhc_fused_backend requires use_fused_mhc=True when set explicitly.") - if self.enable_mhc_connections and self.recompute_granularity == "full": - raise NotImplementedError( - "enable_mhc_connections is not yet compatible with full activation recompute. " - "Use selective recompute with 'mhc' in recompute_modules, or disable " - "activation recompute." - ) - if self.enable_mhc_connections and self.inference_fuse_tp_communication: raise NotImplementedError( "enable_mhc_connections is not compatible with inference_fuse_tp_communication. " diff --git a/tests/unit_tests/models/test_hybrid_mhc.py b/tests/unit_tests/models/test_hybrid_mhc.py index f92c5694c51..2629d6d60b8 100644 --- a/tests/unit_tests/models/test_hybrid_mhc.py +++ b/tests/unit_tests/models/test_hybrid_mhc.py @@ -11,7 +11,10 @@ from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer import TransformerConfig -from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.module import ( + MegatronModule, + convert_module_to_dtype_except_fp32_marked, +) from megatron.core.transformer.moe.moe_layer import MoELayer from megatron.core.transformer.multi_token_prediction import MultiTokenPredictionLayer from megatron.core.transformer.spec_utils import ModuleSpec @@ -160,6 +163,22 @@ def test_fused_backend_policy_is_bound_per_wrapper(self): ): assert getattr(hyper_connection, op_name).keywords["backend"] == "native" + def test_wrapped_residual_layer_does_not_double_count_input(self): + config = _get_config(num_layers=1) + wrapper = _get_stack(config, num_local_layers=1).cuda().layers[0] + hidden = torch.randn( + 8, 2, config.hidden_size * config.mhc_num_residual_streams, device="cuda" + ) + aggregated, h_res, h_post, residual = wrapper.hyper_connection(hidden, return_residual=True) + branch = 0.125 * wrapper.inner_layer.proj(aggregated) + expected = wrapper.hyper_connection.fused_h_res_h_post_bda( + h_res, residual, h_post, (branch, None), dropout_prob=0.0, training=True, fused=False + ) + + actual, _ = wrapper(hidden, attention_mask=None) + + torch.testing.assert_close(actual, expected) + def test_wrapped_gdn_preserves_dynamic_inference_state_shapes(self): config = _get_config(num_layers=1) submodules = _get_dummy_submodules() @@ -184,8 +203,13 @@ def test_wrapped_gdn_preserves_dynamic_inference_state_shapes(self): "recompute_modules": ["core_attn", "mhc"], "mhc_recompute_layer_num": 2, }, + { + "recompute_granularity": "full", + "recompute_method": "uniform", + "recompute_num_layers": 1, + }, ], - ids=["none", "selective_mhc"], + ids=["none", "selective_mhc", "full"], ) def test_forward_backward(self, recompute_kwargs): config = _get_config(num_layers=3, **recompute_kwargs) @@ -208,6 +232,42 @@ def test_forward_backward(self, recompute_kwargs): for name in ("hc_head_fn", "hc_head_base", "hc_head_scale"): assert getattr(stack, name).grad is not None + @pytest.mark.parametrize("recompute_method", ["uniform", "block"]) + @pytest.mark.parametrize("precision", ["fp32", "bf16_fp32_mixing", "bf16_fused"]) + def test_full_recompute_matches_forward_backward(self, recompute_method, precision): + dtype = torch.float32 if precision == "fp32" else torch.bfloat16 + config_kwargs = dict( + num_layers=3, + bf16=dtype == torch.bfloat16, + params_dtype=dtype, + use_fused_mhc=precision == "bf16_fused", + mhc_norm_eps_inside_sqrt=precision == "bf16_fp32_mixing", + mhc_keep_mappings_in_fp32=precision == "bf16_fp32_mixing", + ) + reference = _get_stack(_get_config(**config_kwargs), num_local_layers=3).cuda() + config = _get_config( + **config_kwargs, + recompute_granularity="full", + recompute_method=recompute_method, + recompute_num_layers=2, + ) + checkpointed = _get_stack(config, num_local_layers=3).cuda() + for stack in (reference, checkpointed): + convert_module_to_dtype_except_fp32_marked(stack, dtype) + checkpointed.load_state_dict(reference.state_dict()) + hidden = torch.randn(8, 2, config.hidden_size, device="cuda", dtype=dtype) + reference_input = hidden.clone().requires_grad_() + checkpointed_input = hidden.clone().requires_grad_() + + expected = reference(reference_input, attention_mask=None) + actual = checkpointed(checkpointed_input, attention_mask=None) + torch.testing.assert_close(actual, expected) + expected.float().square().mean().backward() + actual.float().square().mean().backward() + torch.testing.assert_close(checkpointed_input.grad, reference_input.grad) + for name, param in checkpointed.named_parameters(): + torch.testing.assert_close(param.grad, reference.get_parameter(name).grad, msg=name) + def test_fused_bf16_forward_backward(self): config = _get_config( num_layers=2, bf16=True, params_dtype=torch.bfloat16, use_fused_mhc=True diff --git a/tests/unit_tests/transformer/test_hyper_connection_recompute.py b/tests/unit_tests/transformer/test_hyper_connection_recompute.py index 568111f15a6..649b22e0d0f 100644 --- a/tests/unit_tests/transformer/test_hyper_connection_recompute.py +++ b/tests/unit_tests/transformer/test_hyper_connection_recompute.py @@ -460,17 +460,6 @@ def test_config_rejects_fp32_residual_connection(self): @pytest.mark.parametrize( "extra_kwargs, error_type, match", [ - ( - # recompute_num_layers is required for non-selective granularity, and that - # check runs first — supply it so the mHC guard is what actually fires. - { - "recompute_granularity": "full", - "recompute_method": "uniform", - "recompute_num_layers": 1, - }, - NotImplementedError, - "full activation recompute", - ), ({"inference_fuse_tp_communication": True}, NotImplementedError, "single-stream"), ({"mhc_sinkhorn_iterations": 0}, ValueError, "mhc_sinkhorn_iterations"), ({"mhc_init_gating_factor": -0.1}, ValueError, "mhc_init_gating_factor"), diff --git a/tests/unit_tests/transformer/test_mhc_precision.py b/tests/unit_tests/transformer/test_mhc_precision.py new file mode 100644 index 00000000000..d0e75c277d4 --- /dev/null +++ b/tests/unit_tests/transformer/test_mhc_precision.py @@ -0,0 +1,83 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import pytest +import torch + +from megatron.core.transformer.hyper_connection import HyperConnectionModule +from megatron.core.transformer.transformer_config import TransformerConfig + + +@pytest.mark.parametrize("fp32_mixing", [False, True]) +@pytest.mark.parametrize("input_scale", [1.0, 0.001]) +def test_mhc_precision(fp32_mixing, input_scale): + torch.manual_seed(42) + config = TransformerConfig( + num_layers=1, + hidden_size=32, + num_attention_heads=4, + layernorm_epsilon=1e-5, + mhc_norm_eps_inside_sqrt=fp32_mixing, + mhc_keep_mappings_in_fp32=fp32_mixing, + ) + layer = HyperConnectionModule(config, layer_number=1).cuda() + expected_eps = config.layernorm_epsilon if fp32_mixing else 1e-6 + assert layer.norm_eps == expected_eps + with torch.no_grad(): + layer.mapping_proj.weight.normal_(std=0.2) + layer.bias.normal_() + residual = (torch.randn(16, 2, 128, device="cuda") * input_scale).to(torch.bfloat16) + residual.requires_grad_() + output = torch.randn(16, 2, 32, device="cuda", dtype=torch.bfloat16, requires_grad=True) + pre, post, comb = layer.compute_mappings(residual) + expected_dtype = torch.float32 if fp32_mixing else torch.bfloat16 + assert pre.dtype == post.dtype == comb.dtype == expected_dtype + + x = residual.float() + rms = ( + torch.rsqrt(x.square().mean(-1, keepdim=True) + expected_eps) + if fp32_mixing + else (x.norm(dim=-1, keepdim=True) / 128**0.5 + expected_eps).reciprocal() + ) + scales = torch.cat( + [layer.alpha_pre.expand(4), layer.alpha_post.expand(4), layer.alpha_res.expand(16)] + ) + logits = (x @ layer.mapping_proj.weight.T) * rms * scales + layer.bias + expected_pre = (logits[..., :4].sigmoid() + 1e-6).to(expected_dtype) + expected_post = (2 * logits[..., 4:8].sigmoid()).to(expected_dtype) + expected_comb = logits[..., 8:].reshape(16, 2, 4, 4).softmax(-1) + 1e-6 + expected_comb = expected_comb / (expected_comb.sum(-2, keepdim=True) + 1e-6) + for _ in range(config.mhc_sinkhorn_iterations - 1): + expected_comb = expected_comb / (expected_comb.sum(-1, keepdim=True) + 1e-6) + expected_comb = expected_comb / (expected_comb.sum(-2, keepdim=True) + 1e-6) + torch.testing.assert_close(pre, expected_pre) + torch.testing.assert_close(post, expected_post) + torch.testing.assert_close(comb, expected_comb.to(expected_dtype)) + + aggregated = layer.aggregate(residual, pre) + streams = residual.view(16, 2, 4, 32) + expected_aggregate = ( + (streams.float() * expected_pre.float().unsqueeze(-1)).sum(2).to(residual.dtype) + ) + torch.testing.assert_close(aggregated, expected_aggregate) + actual = layer.fused_h_res_h_post_bda(comb, residual, post, (output, None), 0.0, True, False) + expected_mix = torch.einsum("...ij,...ih->...jh", comb, streams.to(expected_dtype)) + expected_output = ( + expected_mix.float() + post.float().unsqueeze(-1) * output.float().unsqueeze(2) + ).to(residual.dtype) + torch.testing.assert_close(actual.view_as(streams), expected_output) + (actual.float().square().mean() + aggregated.float().square().mean()).backward() + for tensor in (residual, output, layer.mapping_proj.weight, layer.bias): + assert tensor.grad is not None and torch.isfinite(tensor.grad).all() + + +def test_mhc_fused_rejects_unsupported_precision(): + config = TransformerConfig( + num_layers=1, + hidden_size=32, + num_attention_heads=4, + enable_mhc_connections=True, + use_fused_mhc=True, + mhc_keep_mappings_in_fp32=True, + ) + with pytest.raises(ValueError, match="Fused mHC does not support"): + HyperConnectionModule(config, layer_number=1) From bdcb18074c0e673488eb3f0be9e6fda7094545dd Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Thu, 10 Sep 2026 01:04:06 +0300 Subject: [PATCH 5/6] KPool DSA Signed-off-by: Hollow Man --- .../absorbed_mla.py | 88 +++- .../experimental_attention_variant/dsa.py | 495 +++++++++++++++--- .../dsa_cudnn_kernels.py | 5 + .../core/transformer/transformer_config.py | 16 + .../test_attention_variant_dsa.py | 22 + .../test_kpool_causal_tail.py | 64 +++ 6 files changed, 604 insertions(+), 86 deletions(-) create mode 100644 tests/unit_tests/transformer/experimental_attention_variant/test_kpool_causal_tail.py diff --git a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py index 9991e6828d1..f477d87850b 100644 --- a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py +++ b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py @@ -20,6 +20,7 @@ from megatron.core import tensor_parallel from megatron.core.extensions.transformer_engine import HAVE_TE +from megatron.core.fp8_utils import get_fp8_disabled_context from megatron.core.models.common.embeddings import ( RotaryEmbedding, YarnRotaryEmbedding, @@ -192,7 +193,11 @@ def __init__( self.cache_mla_latents = self.config.cache_mla_latents assert not self.cache_mla_latents, "cache_mla_latents is not supported for AbsorbedMLA" - if self.config.rope_type == "rope": + # NoPE has no positional slice to embed or rotate. + self.use_rope = self.config.qk_pos_emb_head_dim > 0 + if not self.use_rope: + self.rotary_pos_emb = None + elif self.config.rope_type == "rope": self.rotary_pos_emb = RotaryEmbedding( self.config.qk_pos_emb_head_dim, rotary_percent=self.config.rotary_percent, @@ -416,29 +421,36 @@ def get_query_key_value_tensors( # ========================================= # Prepare RoPE and seqlen related params # ========================================= - rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( - inference_context, None, hidden_states, self.config, packed_seq_params - ) - mscale = 1.0 rotary_pos_cos = None rotary_pos_sin = None packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' - if self.config.rope_type == "rope": - rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq) - else: - if self.config.apply_rope_fusion: - rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( - rotary_seq_len, dtype=hidden_states.dtype, packed_seq=packed_seq - ) - rotary_pos_emb = None - assert inference_context is None, "Inference with MLA RoPE fusion is not supported" - assert ( - fused_apply_mla_rope_for_q is not None - and fused_apply_mla_rope_for_kv is not None - ), "Fused MLA RoPE apply is not imported successfully" + if self.use_rope: + rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( + inference_context, None, hidden_states, self.config, packed_seq_params + ) + if self.config.rope_type == "rope": + rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq) else: - rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq) + if self.config.apply_rope_fusion: + rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( + rotary_seq_len, dtype=hidden_states.dtype, packed_seq=packed_seq + ) + rotary_pos_emb = None + assert ( + inference_context is None + ), "Inference with MLA RoPE fusion is not supported" + assert ( + fused_apply_mla_rope_for_q is not None + and fused_apply_mla_rope_for_kv is not None + ), "Fused MLA RoPE apply is not imported successfully" + else: + rotary_pos_emb, mscale = self.rotary_pos_emb( + rotary_seq_len, packed_seq=packed_seq + ) + else: + # NoPE: no rotary embedding; q_absorbed/kv_compressed carry no pos slice. + rotary_pos_emb = None if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': if packed_seq_params.cu_seqlens_q_padded is not None: @@ -468,7 +480,11 @@ def get_query_key_value_tensors( # q_compressed: [s, b, q_lora_rank / TP] # elif linear_q_down_proj is Linear: # q_compressed: [s / TP, b, q_lora_rank] - q_compressed, _ = self.linear_q_down_proj(hidden_states) + if self.config.mla_disable_attention_fp8: + with get_fp8_disabled_context(self.config): + q_compressed, _ = self.linear_q_down_proj(hidden_states) + else: + q_compressed, _ = self.linear_q_down_proj(hidden_states) # When output is sharded (ColumnParallelLinear), two things are needed to be # identical to a normal Linear. @@ -489,7 +505,11 @@ def get_query_key_value_tensors( # kv_combined: [s, b, (kv_lora_rank + qk_pos_emb_head_dim) / TP] # elif linear_kv_down_proj is Linear: # kv_combined: [s / TP, b, (kv_lora_rank + qk_pos_emb_head_dim)] - kv_combined, _ = self.linear_kv_down_proj(hidden_states) + if self.config.mla_disable_attention_fp8: + with get_fp8_disabled_context(self.config): + kv_combined, _ = self.linear_kv_down_proj(hidden_states) + else: + kv_combined, _ = self.linear_kv_down_proj(hidden_states) if kv_combined.size(-1) != self.config.kv_lora_rank + self.config.qk_pos_emb_head_dim: # kv_combined: [s, b, (kv_lora_rank + qk_pos_emb_head_dim)] kv_combined = gather_from_tensor_model_parallel_region(kv_combined) @@ -547,7 +567,11 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po if self.config.q_lora_rank is not None: # q_compressed: [num_tokens, q_lora_rank] # q: [num_tokens, n * (qk_head_dim + qk_pos_emb_head_dim)] - q, _ = self.linear_q_up_proj(q_compressed) + if self.config.mla_disable_attention_fp8: + with get_fp8_disabled_context(self.config): + q, _ = self.linear_q_up_proj(q_compressed) + else: + q, _ = self.linear_q_up_proj(q_compressed) else: # q_compressed: [num_tokens, hidden_size] # q: [num_tokens, n * (qk_head_dim + qk_pos_emb_head_dim)] @@ -558,11 +582,21 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po # [num_tokens, kv_lora_rank] -> [num_tokens, 1, kv_lora_rank] kv_compressed = torch.unsqueeze(kv_compressed, -2) - # [num_tokens, qk_pos_emb_head_dim] -> [num_tokens, 1, qk_pos_emb_head_dim] - k_pos_emb = torch.unsqueeze(k_pos_emb, -2) k_up_weight, _ = self._get_kv_up_weights() + if not self.use_rope: + q_absorbed = torch.einsum("...nd,ndk->...nk", q, k_up_weight) + q_absorbed = q_absorbed.contiguous() + assert q_absorbed.size(-1) == self.config.kv_lora_rank + assert q_absorbed.is_contiguous() + assert kv_compressed.is_contiguous() + # CheckpointWithoutOutput discards output storage; do not alias its saved input. + return q_absorbed, kv_compressed.clone() + + # [num_tokens, qk_pos_emb_head_dim] -> [num_tokens, 1, qk_pos_emb_head_dim] + k_pos_emb = torch.unsqueeze(k_pos_emb, -2) + if self.config.apply_rope_fusion: # q_no_pe: [num_tokens, n, qk_head_dim] # q_pos_emb: [num_tokens, n, qk_pos_emb_head_dim] @@ -911,7 +945,11 @@ def forward( # ================= # Output. [sq, b, h] # ================= - output, bias = self.linear_proj(core_attn_out) + if self.config.mla_disable_attention_fp8: + with get_fp8_disabled_context(self.config): + output, bias = self.linear_proj(core_attn_out) + else: + output, bias = self.linear_proj(core_attn_out) return output, bias diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index 1d606565d57..dc4da9f08ac 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -6,6 +6,7 @@ from typing import Optional, Tuple, Union import torch +import torch.nn as nn from megatron.core.fp8_utils import get_fp8_disabled_context from megatron.core.models.common.embeddings import ( @@ -664,23 +665,25 @@ def _compute_index_scores( Returns: index_scores: FP32 [batch, seqlen_q, seqlen_k], the index scores. """ - # Compute attention scores: q @ k^T - # [seqlen_q, batch, index_n_heads, index_head_dim] @ [seqlen_k, batch, index_head_dim]^T - # -> [seqlen_q, batch, index_n_heads, seqlen_k] - index_scores = torch.einsum('sbhd,tbd->sbht', q.float(), k.float()) - - # Optionally apply ReLU activation (used by DeepSeek V3.2, not GLM5). - if use_relu: - index_scores = torch.relu(index_scores) - - # Weight each head by attention weights. - # [seqlen_q, batch, index_n_heads, seqlen_k] * [seqlen_q, batch, index_n_heads, 1] - # -> [seqlen_q, batch, index_n_heads, seqlen_k] - index_scores = index_scores * weights.unsqueeze(-1) - - # Sum across attention heads. - # [seqlen_q, batch, index_n_heads, seqlen_k] -> [seqlen_q, batch, seqlen_k] - index_scores = index_scores.sum(dim=2) + sq, batch, n_heads, head_dim = q.shape + sk = k.size(0) + k_fp32 = k.float() + + # Chunk over seqlen_q to avoid materializing the full [sq, batch, heads, sk] + # fp32 tensor. Target ~1 GB per chunk. + bytes_per_token = batch * n_heads * sk * 4 + chunk_size = min(sq, max(1, 1024 * 1024 * 1024 // max(1, bytes_per_token))) + index_scores = torch.empty(sq, batch, sk, dtype=torch.float32, device=q.device) + + for start in range(0, sq, chunk_size): + end = min(start + chunk_size, sq) + # [chunk, batch, heads, sk] + scores = torch.einsum('sbhd,tbd->sbht', q[start:end].float(), k_fp32) + if use_relu: + scores.relu_() + # Weight and sum over heads in one step: [chunk, batch, sk] + index_scores[start:end] = (scores * weights[start:end].unsqueeze(-1)).sum(dim=2) + del scores # Transpose to [batch, seqlen_q, seqlen_k]. index_scores = index_scores.transpose(0, 1) @@ -737,6 +740,257 @@ def fused_qk_topk_naive( return index_scores, topk_indices +def _kpool_fp8_input(x: torch.Tensor) -> torch.Tensor: + """Match the indexer's FP32 Hadamard, BF16 rounding, and E4M3 power-of-two scale.""" + if not x.numel(): + return x.float() + assert hadamard_transform is not None, "fast_hadamard_transform is required for FP8 KPool." + x = hadamard_transform(x.float(), scale=x.shape[-1] ** -0.5).to(torch.bfloat16).float() + absmax = x.abs().amax(dim=-1, keepdim=True).clamp_min(1e-4) + scale = torch.exp2(torch.ceil(torch.log2(absmax / 448.0))) + return (x / scale).clamp(-448, 448).to(torch.float8_e4m3fn).float() * scale + + +def _kpool_compress_keys( + k: torch.Tensor, gate_score: torch.Tensor, ape: torch.Tensor, pool_size: int +) -> torch.Tensor: + """Softmax-weighted pool keys, accumulated in FP32 and returned in BF16. + + Keys and gates are [tokens, batch, head_dim]; ape is [pool_size, head_dim]. + Only complete pools are compressed. Query-local tails are appended separately. + """ + seqlen, bsz, head_dim = k.shape + assert head_dim == ape.shape[1], f"head_dim {head_dim} != ape dim1 {ape.shape[1]}" + num_pools = seqlen // pool_size + # Drop the trailing incomplete pool from compression; its tokens are appended + # later via append_tail_to_topk (always_select_tail). + usable = num_pools * pool_size + # [num_pools, pool_size, batch, head_dim] + k_p = k[:usable].reshape(num_pools, pool_size, bsz, head_dim) + # gate_score: [seqlen, batch, head_dim] -> [num_pools, pool_size, batch, head_dim] + if gate_score is not None: + g = gate_score[:usable].reshape(num_pools, pool_size, bsz, head_dim).float() + else: + g = torch.zeros((num_pools, pool_size, bsz, head_dim), dtype=torch.float32, device=k.device) + + # Per-dim softmax across the pool's slots: score[slot] = gate_score[slot] + ape[slot]. + # ape: [pool_size, head_dim] -> broadcast over (num_pools, batch). + ape_f = ape.to(dtype=torch.float32, device=k.device) # [pool_size, head_dim] + score = g + ape_f.unsqueeze(0).unsqueeze(2) # [num_pools, pool_size, batch, head_dim] + # Numerically-stable per-dim softmax over dim=1 (the pool slot dim). + score_max = score.max(dim=1, keepdim=True).values + prob = torch.exp(score - score_max) + # weighted sum of k over pool slots: [num_pools, batch, head_dim]. Keep the + # numerator 4D ([num_pools, 1, batch, head_dim]) so it broadcasts cleanly + # against the 4D denom ([num_pools, 1, 1, head_dim]); a 3D numerator would + # left-pad and produce a spurious extra (num_pools) dimension. + k_f = k_p.float() + k_pooled = (prob * k_f).sum(dim=1, keepdim=True) / prob.sum(dim=1, keepdim=True).clamp( + min=1e-12 + ) + k_pooled = k_pooled.squeeze(1) # [num_pools, batch, head_dim] + return k_pooled.to(torch.bfloat16) + + +def _expand_pools_to_tokens( + pool_ids: torch.Tensor, + pool_valid: torch.Tensor, + topk_tokens: int, + pool_size: int, + pool_token_base: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Expand fixed-width pool IDs to token IDs, preserving -1 padding.""" + assert pool_ids.ndim == 2 and pool_valid.shape == pool_ids.shape + assert pool_ids.shape[1] * pool_size == topk_tokens + if pool_token_base is None: + starts = pool_ids * pool_size + elif pool_token_base.numel(): + starts = pool_token_base[pool_ids.clamp(min=0)] + else: + starts = torch.zeros_like(pool_ids) + offsets = torch.arange(pool_size, device=pool_ids.device) + tokens = starts.unsqueeze(-1) + offsets + tokens = tokens.masked_fill(~pool_valid.unsqueeze(-1), -1) + return tokens.reshape(pool_ids.shape[0], topk_tokens).to(torch.int32) + + +def _append_tail_to_topk( + topk_result: torch.Tensor, + seq_lens: torch.Tensor, + pool_size: int, + tail_start_override: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Append each query's incomplete causal pool, in global token coordinates.""" + tail_count = seq_lens.to(torch.int32).remainder(pool_size) + tail_start = ( + seq_lens.to(torch.int32) - tail_count + if tail_start_override is None + else tail_start_override.to(torch.int32) + ) + offsets = torch.arange(pool_size - 1, device=topk_result.device) + tail = tail_start[:, None] + offsets + tail = tail.masked_fill(offsets >= tail_count[:, None], -1).to(topk_result.dtype) + return torch.cat((topk_result, tail), dim=-1) + + +def _kpool_compress_keys_per_seg( + k: torch.Tensor, + gate_score: Optional[torch.Tensor], + ape: torch.Tensor, + pool_size: int, + cu_seqlens_kv: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Compress complete pools within each packed segment. + + Return pooled keys and their global starting token indices. Segment boundaries + need not be multiples of pool_size; a pool must never span two documents. + """ + cu = cu_seqlens_kv.to(device=k.device, dtype=torch.int64) + n_seg = int(cu.numel()) - 1 + pooled_parts = [] + base_parts = [] + for i in range(n_seg): + s = int(cu[i]) + e = int(cu[i + 1]) + seg_len = e - s + if seg_len <= 0: + continue + k_seg = k[s:e] + gate_seg = gate_score[s:e] if gate_score is not None else None + k_pooled_seg = _kpool_compress_keys(k_seg, gate_seg, ape, pool_size) + # [num_pools_seg, b, d] + n_pools_seg = k_pooled_seg.size(0) + pooled_parts.append(k_pooled_seg) + # pool j (local) of this segment starts at global token s + j*pool_size. + seg_bases = torch.arange(n_pools_seg, device=k.device, dtype=torch.int64) * pool_size + s + base_parts.append(seg_bases) + k_pooled_global = torch.cat(pooled_parts, dim=0) + pool_token_base = torch.cat(base_parts, dim=0) + return k_pooled_global, pool_token_base + + +def fused_qk_topk_kpool( + q: torch.Tensor, + k: torch.Tensor, + weights: torch.Tensor, + index_topk: int, + pool_size: int, + gate_score: torch.Tensor, + ape: torch.Tensor, + mask: Optional[torch.Tensor] = None, + varlen_starts: Optional[torch.Tensor] = None, + varlen_ends: Optional[torch.Tensor] = None, + key_positions: Optional[torch.Tensor] = None, + cu_seqlens_kv: Optional[torch.Tensor] = None, + use_relu: bool = True, + always_select_tail: bool = True, + fp8_indexer: bool = False, +): + """Select complete causal pools and append each query's incomplete tail. + + q is [queries, batch, heads, dim], k/gate_score are [keys, batch, dim], + and weights are [queries, batch, heads]. Packed bounds use global token + coordinates. Output indices are [batch, queries, index_topk + pool_size - 1] + when always_select_tail is enabled, with -1 for unused slots. + """ + sk = k.size(0) + + # Packed pools restart at each document boundary. + use_per_seg = cu_seqlens_kv is not None and cu_seqlens_kv.numel() >= 2 + if use_per_seg: + k_pooled, pool_token_base = _kpool_compress_keys_per_seg( + k, gate_score, ape, pool_size, cu_seqlens_kv + ) + num_pools = k_pooled.size(0) + else: + num_pools = sk // pool_size + k_pooled = _kpool_compress_keys(k, gate_score, ape, pool_size) + pool_token_base = torch.arange(num_pools, device=k.device, dtype=torch.int64) * pool_size + + if fp8_indexer: + q, k_pooled = _kpool_fp8_input(q), _kpool_fp8_input(k_pooled) + index_scores = _compute_index_scores(q, weights, k_pooled, use_relu=use_relu) + + # A pool is causal only when its final token is within the query's bounds. + pool_positions = pool_token_base + (pool_size - 1) + eff_key_positions = ( + key_positions[pool_positions] if key_positions is not None else pool_positions + ) + v_starts, v_ends, k_pos = dsa_masking.normalize_varlen_bounds( + mask=mask, + varlen_starts=varlen_starts, + varlen_ends=varlen_ends, + key_positions=eff_key_positions, + sk=num_pools, + device=index_scores.device, + ) + if v_starts is not None: + index_scores = dsa_masking.apply_starts_ends_mask_to_scores( + index_scores, v_starts, v_ends, k_pos + ) + elif mask is not None: + assert mask.dtype == index_scores.dtype, "Mask dtype must match index scores dtype" + index_scores = index_scores + mask + + # Keep the selection width fixed, including when fewer causal pools exist. + budget = index_topk // pool_size + select_k = min(budget, num_pools) + if select_k > 0: + topk_scores, pool_topk = index_scores.topk(select_k, dim=-1) + # [batch, seqlen_q, select_k] -> mask invalid pools + pool_topk = pool_topk.masked_fill(topk_scores == float("-inf"), -1) + else: + pool_topk = torch.empty( + index_scores.shape[:-1] + (0,), dtype=torch.int64, device=index_scores.device + ) + if pool_topk.shape[-1] < budget: + pad = torch.full( + index_scores.shape[:-1] + (budget - pool_topk.shape[-1],), + -1, + dtype=torch.int64, + device=index_scores.device, + ) + pool_topk = torch.cat([pool_topk, pad], dim=-1) + + # Expand [batch * queries, pools] to a fixed token budget. + rows = pool_topk.shape[0] * pool_topk.shape[1] + pool_flat = pool_topk.reshape(rows, -1) + pool_valid = pool_flat >= 0 + # Clamp invalid ids to 0 for the arithmetic, restore -1 via the where mask. + safe_pool = pool_flat.clamp(min=0) + token_topk = _expand_pools_to_tokens( + safe_pool, + pool_valid, + index_topk, + pool_size, + pool_token_base=pool_token_base if use_per_seg else None, + ) + # token_topk is [rows, index_topk]; reshape back to [batch, seqlen_q, index_topk]. + token_topk = token_topk.reshape(pool_topk.shape[0], pool_topk.shape[1], index_topk) + + if always_select_tail: + # Pool phase is query-local, never the final length of the packed sample. + sq, batch = q.shape[:2] + ends = v_ends if v_ends is not None else torch.arange(1, sq + 1, device=q.device) + if v_starts is not None: + starts = v_starts + elif cu_seqlens_kv is not None: + cu = cu_seqlens_kv.to(device=q.device, dtype=torch.int64) + starts = cu[torch.searchsorted(cu[1:], ends - 1, right=True)] + else: + starts = torch.zeros_like(ends) + lengths = ends - starts + tail_starts = ends - lengths.remainder(pool_size) + token_topk = _append_tail_to_topk( + token_topk.reshape(rows, -1), + lengths.expand(batch, -1).reshape(-1), + pool_size, + tail_start_override=tail_starts.expand(batch, -1).reshape(-1), + ).reshape(batch, sq, -1) + + return index_scores, token_topk + + def fwd_fused_indexer_loss_naive( q, weights, @@ -972,36 +1226,47 @@ def bwd_fused_indexer_loss_naive( grad_weighted_scores = grad_index_scores.unsqueeze(2) # [sq, b, 1, sk] del grad_index_scores - # Compute forward values needed for backward - scores = torch.einsum('sbhd,tbd->sbht', q.float(), k.float()) # [sq, b, h, sk] - - # Backward through multiplication by weights (with optional ReLU). - if use_relu: - scores_for_weights = torch.relu(scores) - relu_mask = scores > 0 - else: - scores_for_weights = scores - relu_mask = None - del scores + # Chunk over seqlen_q to avoid materializing the full [sq, b, h, sk] fp32 tensor. + sq_q, b_q, h_q, d_q = q.shape + k_fp32 = k.float() + bytes_per_token = b_q * h_q * sk * 4 + chunk_size = min(sq_q, max(1, 1024 * 1024 * 1024 // max(1, bytes_per_token))) + + grad_q = torch.empty(sq_q, b_q, h_q, d_q, dtype=torch.float32, device=q.device) + grad_weights = torch.empty(sq_q, b_q, h_q, dtype=torch.float32, device=q.device) + grad_k = torch.zeros(sk, b_q, d_q, dtype=torch.float32, device=q.device) + + for start in range(0, sq_q, chunk_size): + end = min(start + chunk_size, sq_q) + q_chunk = q[start:end] # [chunk, b, h, d] + gw_chunk = grad_weighted_scores[start:end] # [chunk, b, 1, sk] + w_chunk = weights[start:end] # [chunk, b, h] + + # Forward scores for this chunk: [chunk, b, h, sk] + scores_chunk = torch.einsum('sbhd,tbd->sbht', q_chunk.float(), k_fp32) + if use_relu: + relu_mask_chunk = scores_chunk > 0 + scores_chunk.relu_() + else: + relu_mask_chunk = None - # ∂L/∂weights = grad * scores_for_weights (sum over sk) - grad_weights = (grad_weighted_scores * scores_for_weights).sum(dim=-1) # [sq, b, h] + # ∂L/∂weights = grad * scores (sum over sk) + grad_weights[start:end] = (gw_chunk * scores_chunk).sum(dim=-1) + del scores_chunk - # ∂L/∂scores = grad * weights - grad_scores = grad_weighted_scores * weights.unsqueeze(-1) # [sq, b, h, sk] - del grad_weighted_scores, scores_for_weights + # ∂L/∂scores = grad * weights: [chunk, b, h, sk] + grad_scores_chunk = gw_chunk * w_chunk.unsqueeze(-1) + if use_relu: + grad_scores_chunk.masked_fill_(~relu_mask_chunk, 0.0) + del relu_mask_chunk - # Backward through ReLU (skip when use_relu=False) - if use_relu: - grad_scores = grad_scores * relu_mask.float() - del relu_mask + # ∂L/∂q = einsum('sbht,tbd->sbhd', grad_scores, k) + grad_q[start:end] = torch.einsum('sbht,tbd->sbhd', grad_scores_chunk, k_fp32) + # ∂L/∂k = einsum('sbht,sbhd->tbd', grad_scores, q) (accumulate) + grad_k += torch.einsum('sbht,sbhd->tbd', grad_scores_chunk, q_chunk.float()) + del grad_scores_chunk - # Backward through einsum 'sbhd,tbd->sbht' - # ∂L/∂q = einsum('sbht,tbd->sbhd', grad_scores, k) - grad_q = torch.einsum('sbht,tbd->sbhd', grad_scores, k.float()) # [sq, b, h, d] - # ∂L/∂k = einsum('sbht,sbhd->tbd', grad_scores, q) - grad_k = torch.einsum('sbht,sbhd->tbd', grad_scores, q.float()) # [sk, b, d] - del grad_scores + del grad_weighted_scores return grad_q.to(q.dtype), grad_weights.to(weights.dtype), grad_k.to(k.dtype) @@ -1279,7 +1544,13 @@ def __init__( self.pg_collection = pg_collection # Initialize Position Embedding. - if self.config.rope_type == 'rope': + # NoPE (qk_pos_emb_head_dim == 0, e.g. GLM-5.3-Flash indexer): skip the + # rotary embedding entirely; constructing it with dim=0 yields empty + # inv_freq / NaN, and the forward path skips RoPE split/apply. + self.use_rope = self.qk_pos_emb_head_dim > 0 + if not self.use_rope: + self.rotary_pos_emb = None + elif self.config.rope_type == 'rope': self.rotary_pos_emb = RotaryEmbedding( self.qk_pos_emb_head_dim, rotary_percent=self.config.rotary_percent, @@ -1355,6 +1626,27 @@ def __init__( skip_weight_param_allocation=False, parallel_mode="duplicated", ) + + # Pool compression uses replicated per-token gates and slot-position biases. + self.index_kpool = int(self.config.dsa_indexer_kpool) + self.index_kpool_always_select_tail = bool(self.config.dsa_indexer_kpool_always_select_tail) + # Per-token gate score for the kpool path; set in forward_before_topk. + self._kpool_gate_score: Optional[torch.Tensor] = None + if self.index_kpool > 1: + # fp32 [kpool, index_head_dim] additive positional bias per pool slot. + self.index_kpool_compress_ape = torch.nn.Parameter( + torch.zeros(self.index_kpool, self.index_head_dim, dtype=torch.float32) + ) + # bf16 [index_head_dim, hidden_size]; gate_score = F.linear(x, gate) = x @ gate^T + # -> [seqlen, index_head_dim]. Matches vLLM's checkpoint name (no .weight suffix). + self.index_kpool_compress_gate = torch.nn.Parameter( + torch.empty(self.index_head_dim, self.hidden_size, dtype=torch.bfloat16) + ) + nn.init.normal_(self.index_kpool_compress_gate, std=0.01) + else: + self.index_kpool_compress_ape = None + self.index_kpool_compress_gate = None + # Indexer projections are duplicated across tensor-parallel ranks, so their gradients # should be averaged during final gradient synchronization. for param in self.parameters(): @@ -1368,6 +1660,9 @@ def _apply_rope( cu_seqlens: Optional[torch.Tensor] = None, ): """Apply RoPE to the input tensor.""" + # NoPE: no positional component; x is all nope, return unchanged. + if not self.use_rope: + return x # x_pe [seqlen, batch, *, qk_pos_emb_head_dim] # x_nope [seqlen, batch, *, index_head_dim - qk_pos_emb_head_dim] # To align with DeepSeek's implementation, @@ -1407,14 +1702,19 @@ def forward_before_topk( # ========================================= # Prepare RoPE params # ========================================= - rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( - None, None, x, self.config, packed_seq_params - ) - if self.config.rope_type == "rope": - rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq) - mscale = 1.0 + if self.use_rope: + rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( + None, None, x, self.config, packed_seq_params + ) + if self.config.rope_type == "rope": + rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq) + mscale = 1.0 + else: + rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq) else: - rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq) + # NoPE: no rotary embedding; _apply_rope is a no-op. + rotary_pos_emb = None + mscale = 1.0 if packed_seq: cu_seqlens_q, cu_seqlens_kv = dsa_layout.get_packed_qk_cu_seqlens(packed_seq_params) else: @@ -1442,7 +1742,9 @@ def forward_before_topk( # -> [seqlen, batch, index_n_heads, index_head_dim] q = q.reshape(seqlen, bsz, self.index_n_heads, self.index_head_dim) q = self._apply_rope(q, rotary_pos_emb, mscale, cu_seqlens=cu_seqlens_q) - if self.config.dsa_indexer_rotate_activation: + if self.config.dsa_indexer_rotate_activation and not ( + self.index_kpool > 1 and self.config.dsa_indexer_kpool_fp8 + ): q = rotate_activation(q) # ========================================= @@ -1463,18 +1765,34 @@ def forward_before_topk( k = k.reshape(seqlen, bsz, self.index_head_dim) # ========================================= - # Rotate activation (k only; q already rotated in FP8 path) + # Rotate activation (k only; q already rotated in FP8 path). + # Pooled keys are rotated after compression by _kpool_fp8_input. # ========================================= - if self.config.dsa_indexer_rotate_activation: + if self.config.dsa_indexer_rotate_activation and self.index_kpool <= 1: k = rotate_activation(k) # ========================================= # Prepare weights for index scores # ========================================= # [seqlen, batch, hidden_size] -> [seqlen, batch, index_n_heads] - weights, _ = self.linear_weights_proj(x) + # The pooled indexer keeps the head-gate projection in FP32. + if self.index_kpool > 1: + weights = torch.nn.functional.linear( + x.float(), self.linear_weights_proj.weight.float() + ) + else: + weights, _ = self.linear_weights_proj(x) weights = weights * (self.index_n_heads**-0.5) * self.softmax_scale + # Save token-aligned compression scores for the subsequent pool selection. + if self.index_kpool > 1 and self.index_kpool_compress_gate is not None: + with get_fp8_disabled_context(self.config): + self._kpool_gate_score = torch.nn.functional.linear( + x, self.index_kpool_compress_gate + ) + else: + self._kpool_gate_score = None + return q, k, weights def forward_with_scores( @@ -1505,10 +1823,30 @@ def forward_with_scores( # [seqlen, batch, index_n_heads] q, k, weights = self.forward_before_topk(x, qr, packed_seq_params) - # [batch, seqlen, seqlen], [batch, seqlen, index_topk] - index_scores, topk_indices = fused_qk_topk_naive( - q, k, weights, self.index_topk, mask, use_relu=self.config.dsa_indexer_scoring_relu - ) + if self.index_kpool > 1 and self._kpool_gate_score is not None: + # Select pools, then expand them to token indices. + _cu_kv = None + if packed_seq_params is not None and packed_seq_params.qkv_format == "thd": + _cu_kv, _ = dsa_layout.get_packed_qk_cu_seqlens(packed_seq_params) + index_scores, topk_indices = fused_qk_topk_kpool( + q, + k, + weights, + self.index_topk, + self.index_kpool, + self._kpool_gate_score, + self.index_kpool_compress_ape, + mask=mask, + cu_seqlens_kv=_cu_kv, + use_relu=self.config.dsa_indexer_scoring_relu, + always_select_tail=self.index_kpool_always_select_tail, + fp8_indexer=self.config.dsa_indexer_kpool_fp8, + ) + else: + # [batch, seqlen, seqlen], [batch, seqlen, index_topk] + index_scores, topk_indices = fused_qk_topk_naive( + q, k, weights, self.index_topk, mask, use_relu=self.config.dsa_indexer_scoring_relu + ) return index_scores, topk_indices @@ -2144,6 +2482,14 @@ def _build_kv_reorder_idx(local_len): f"k_seqlen={k.size(0)}, expected={kv_reorder_idx.numel()}" ) k = k.index_select(0, kv_reorder_idx) + # Apply the same CP gather + reorder to the kpool gate score so it + # matches the now-global key length (otherwise the per-seg reshape + # in _kpool_compress_keys crashes with a size mismatch). + gate = self.indexer._kpool_gate_score + if gate is not None and gate.size(0) in local_cp_kv_lens: + gate = gather_from_sequence_parallel_region(gate, group=cp_group) + gate = gate.index_select(0, kv_reorder_idx) + self.indexer._kpool_gate_score = gate if sequence_parallel_tp and q.size(0) != sq: if ( q.size(0) != sequence_parallel_tp_full_rows @@ -2190,7 +2536,12 @@ def compute_indexer_loss_with_reference_path(): ) fused_output = None - if use_fused_kernels and not self.index_share: + # kpool DSA indexer (GLM-5.3-Flash, index_kpool > 1): the fused cuDNN DSA + # path computes its own per-token top-k and does NOT implement pool-granular + # selection / key compression / tail-append. Bypass it so the Python kpool + # top-k path below runs instead (mirrors DSAIndexer.forward_with_scores). + _is_kpool = self.indexer is not None and getattr(self.indexer, "index_kpool", 1) > 1 + if use_fused_kernels and not self.index_share and not _is_kpool: assert q is not None and k is not None and weights is not None fused_output = dsa_kernels.run_fused_dsa_attention( config=self.config, @@ -2327,7 +2678,29 @@ def slice_topk_to_local_sequence_parallel_rows(): # =================================== # Get top-k indices # =================================== - if fused_bounds is not None: + if _is_kpool: + # KPool selection is discrete and does not need an autograd graph. + with torch.no_grad(): + _index_scores, topk_indices = fused_qk_topk_kpool( + q, + k, + weights, + self.index_topk, + self.indexer.index_kpool, + self.indexer._kpool_gate_score, + self.indexer.index_kpool_compress_ape, + mask=float_mask, + varlen_starts=varlen_starts, + varlen_ends=varlen_ends, + key_positions=key_positions, + cu_seqlens_kv=cu_seqlens_kv if packed_thd else None, + use_relu=self.config.dsa_indexer_scoring_relu, + always_select_tail=self.indexer.index_kpool_always_select_tail, + fp8_indexer=self.config.dsa_indexer_kpool_fp8, + ) + del _index_scores + slice_topk_to_local_sequence_parallel_rows() + elif fused_bounds is not None: starts_i32, ends_i32 = fused_bounds block_size = int(getattr(self, "fused_indexer_block_size", 8192)) fused_topk = dsa_kernels.run_fused_qk_topk( diff --git a/megatron/core/transformer/experimental_attention_variant/dsa_cudnn_kernels.py b/megatron/core/transformer/experimental_attention_variant/dsa_cudnn_kernels.py index 640db5c91ec..824d608886f 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa_cudnn_kernels.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa_cudnn_kernels.py @@ -1432,6 +1432,11 @@ def run_fused_qk_topk_with_loss( use_local_indexer_varlen, packed_seq_params, single_packed_thd_sequence, cp_size ) ) + # The kpool indexer path (index_kpool > 1) computes weights in fp32 for + # numerical accuracy, but the cuDNN bf16 indexer kernel requires bfloat16. + # Cast to match q's dtype (always bf16 under get_fp8_disabled_context). + if weights.dtype != q.dtype: + weights = weights.to(dtype=q.dtype) return FusedQKTopKWithSparseLossFunc.apply( q.contiguous(), k.contiguous(), diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 2a3e86ed2b1..dab67242487 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -365,12 +365,28 @@ class TransformerConfig(ModelParallelConfig): dsa_indexer_scoring_relu: bool = True """Whether DSA indexer should apply ReLU to q@k^T scores before weighting.""" + dsa_indexer_kpool_fp8: bool = False + """Match FP8 KPool index scores + model weights and attention remain in their configured dtype.""" + dsa_indexer_k_norm_epsilon: Optional[float] = None """Optional epsilon override for the DSA indexer key LayerNorm.""" dsa_indexer_k_norm_fp32: bool = False """Whether DSA indexer key LayerNorm should run on fp32 inputs.""" + mla_disable_attention_fp8: bool = False + """Force MLA attention GEMMs (q_a_proj, q_b_proj, kv_a_proj, o_proj) to BF16 + even under FP8 training. This aligns the actor's attention path with vLLM + rollout, which runs the main MLA attention in BF16. Only the attention GEMMs + (~2.5% of model params) are affected; MoE experts stay FP8.""" + + dsa_indexer_kpool: int = 1 + """Number of keys per softmax-weighted indexer pool; 1 keeps per-token selection.""" + + dsa_indexer_kpool_always_select_tail: bool = False + """Append each query's incomplete causal pool after the selected history tokens.""" + #################### # Compressed sparse attention #################### diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py index f7f190fbc62..321d0a39b39 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py @@ -2401,6 +2401,28 @@ def test_dsa_indexer_constructor(self, seqlen): assert self.indexer.index_topk == 32 assert self.indexer.k_norm.eps == pytest.approx(1e-6) + def test_kpool_projection_precision_and_backward(self, seqlen): + self.indexer.cuda() + x = torch.randn(seqlen, 1, 256, device="cuda", dtype=torch.bfloat16, requires_grad=True) + qr = torch.randn(seqlen, 1, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True) + gate = torch.nn.Parameter(torch.randn(64, 256, device="cuda", dtype=torch.bfloat16)) + with ( + patch.object(self.indexer, "index_kpool", 4), + patch.object(self.indexer, "index_kpool_compress_gate", gate), + patch.object(self.config, "dsa_indexer_rotate_activation", False), + ): + q, k, weights = self.indexer.forward_before_topk(x, qr) + gate_score = self.indexer._kpool_gate_score + assert q.dtype == k.dtype == gate_score.dtype == torch.bfloat16 + assert weights.dtype == torch.float32 + torch.testing.assert_close(gate_score, torch.nn.functional.linear(x, gate)) + ( + q.float().sum() + k.float().sum() + weights.sum() + gate_score.float().sum() + ).backward() + assert torch.isfinite(x.grad).all() + assert torch.isfinite(qr.grad).all() + assert torch.isfinite(gate.grad).all() + @pytest.mark.parametrize("interleaved", [False, True]) def test_dsa_indexer_rope_interleave_follows_config(self, seqlen, interleaved): """Ensure indexer RoPE uses the model-configured interleave convention.""" diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_kpool_causal_tail.py b/tests/unit_tests/transformer/experimental_attention_variant/test_kpool_causal_tail.py new file mode 100644 index 00000000000..e64cc43a6c1 --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_kpool_causal_tail.py @@ -0,0 +1,64 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import pytest +import torch + +from megatron.core.transformer.experimental_attention_variant.dsa import ( + _kpool_fp8_input, + fused_qk_topk_kpool, +) +from megatron.core.transformer.experimental_attention_variant.dsa_masking import ( + generate_varlen_mask_params_for_positions, +) + + +@pytest.mark.parametrize("lengths", [(8,), (7,), (3,), (3, 6), (5, 9, 3)]) +@pytest.mark.parametrize("query_stride", [1, 2]) +@pytest.mark.parametrize("explicit_key_positions", [False, True]) +@pytest.mark.parametrize("fp8_indexer", [False, True]) +def test_kpool_preserves_each_query_causal_tail( + lengths, query_stride, explicit_key_positions, fp8_indexer +): + torch.manual_seed(123) + cu = torch.tensor([0, *torch.tensor(lengths).cumsum(0).tolist()], device="cuda") + total = sum(lengths) + positions = torch.arange(0, total, query_stride, device="cuda") + starts, ends = generate_varlen_mask_params_for_positions(cu, positions) + q = torch.randn(len(positions), 1, 2, 8, device="cuda") + k = torch.randn(total, 1, 8, device="cuda") + weights = torch.ones(len(positions), 1, 2, device="cuda") + _, indices = fused_qk_topk_kpool( + q, + k, + weights, + index_topk=16, + pool_size=4, + gate_score=torch.zeros_like(k), + ape=torch.zeros(4, 8, device="cuda"), + varlen_starts=starts, + varlen_ends=ends, + key_positions=torch.arange(total, device="cuda") if explicit_key_positions else None, + cu_seqlens_kv=cu, + fp8_indexer=fp8_indexer, + ) + assert indices.shape == (1, len(positions), 19) + for row, start, end in zip(indices[0], starts.tolist(), ends.tolist()): + # Below the pool budget, sparse attention must contain the full causal prefix. + actual = row[row >= 0].sort().values + expected = torch.arange(start, end, device="cuda", dtype=actual.dtype) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("input_scale", [0.0, 1e-5, 1.0, 1000.0]) +def test_kpool_fp8_input_matches_hadamard_matrix_reference(input_scale): + torch.manual_seed(456) + x = (torch.randn(33, 128, device="cuda") * input_scale).to(torch.bfloat16) + matrix = torch.ones(1, 1, device="cuda") + for _ in range(7): + matrix = torch.cat((torch.cat((matrix, matrix), 1), torch.cat((matrix, -matrix), 1)), 0) + rotated = (x.float() @ matrix / 128**0.5).to(torch.bfloat16).float() + scale = torch.exp2( + torch.ceil(torch.log2(rotated.abs().amax(-1, keepdim=True).clamp_min(1e-4) / 448)) + ) + expected = (rotated / scale).to(torch.float8_e4m3fn).float() * scale + torch.testing.assert_close(_kpool_fp8_input(x), expected, rtol=0, atol=0) From 3fceb07159d3ebc2a4fe175a2659e158d3a3df25 Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Sat, 12 Sep 2026 01:01:21 +0300 Subject: [PATCH 6/6] Guard backward (E5M2) wgrad against all-zero gradient blocks Signed-off-by: Hollow Man --- megatron/core/fp8_utils.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py index 80fcd41f3ba..9cd196cb336 100644 --- a/megatron/core/fp8_utils.py +++ b/megatron/core/fp8_utils.py @@ -779,8 +779,14 @@ def get_fp8_recipe(config: TransformerConfig): fp8_format=fp8_format, fp8_dpa=config.fp8_dot_product_attention ) elif config.fp8_recipe == Fp8Recipe.blockwise and is_te_min_version("2.3.0.dev0"): - fp8_recipe = transformer_engine.common.recipe.Float8BlockScaling( - fp8_format=fp8_format + # Guard backward (E5M2) wgrad against all-zero gradient blocks: + # amax=0 → scale=inf → NaN. + _cls = transformer_engine.common.recipe.Float8BlockScaling + fp8_recipe = _cls( + fp8_format=fp8_format, + fp8_quant_bwd_grad=transformer_engine.common.recipe.QParams( + power_2_scale=_cls.fp8_quant_bwd_grad.power_2_scale, amax_epsilon=1e-12 + ), ) elif config.fp8_recipe == Fp8Recipe.mxfp8: fp8_recipe = transformer_engine.common.recipe.MXFP8BlockScaling(