diff --git a/megatron/core/models/common/embeddings/rope_utils.py b/megatron/core/models/common/embeddings/rope_utils.py
index b990615da29..e104513057f 100644
--- a/megatron/core/models/common/embeddings/rope_utils.py
+++ b/megatron/core/models/common/embeddings/rope_utils.py
@@ -220,49 +220,49 @@ def _apply_rotary_pos_emb_thd(
     cp_size = cp_group.size()
     cp_rank = cp_group.rank()
     seqlens = ((cu_seqlens[1:] - cu_seqlens[:-1]) // cp_size).tolist()
+    sequence_splits = torch.split(t, seqlens)
+    total_seqlen = int(cu_seqlens[-1].item())
+    has_packed_freqs = freqs.dim() >= 1 and freqs.size(0) == total_seqlen
 
     # Handle two different frequency tensor formats:
     # 1. If freqs.size(0) == cu_seqlens[-1]: freqs contains all positions across all sequences
     #    -> Use offset-based mapping for exact positional correspondence
     # 2. Otherwise: freqs contains only max sequence length positions
     #    -> Use traditional mapping without offsets (map first :seqlen part)
-    if freqs.dim() >= 1 and freqs.size(0) == cu_seqlens[-1]:
+    if has_packed_freqs:
         # CASE 1: Exact mapping with offsets
-        # Build packed freqs in one pass, then apply once to the whole packed tensor
-        sequence_splits = torch.split(t, seqlens)
-        freq_slices = []
+        local_freqs = []
         for i, x in enumerate(sequence_splits):
             # cu_seqlens[i] is the starting offset of this sequence in the original batch
             seq_start_offset = cu_seqlens[i].item()
-            freq_slices.append(
+            local_freqs.append(
                 _get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs, seq_start_offset)
             )
-
-        freqs_packed = torch.cat(freq_slices, dim=0)
-
+        freqs = torch.cat(local_freqs, dim=0)
         return _apply_rotary_pos_emb_bshd(
             t.unsqueeze(1),
-            freqs_packed,
+            freqs,
             rotary_interleaved=rotary_interleaved,
             mla_rotary_interleaved=mla_rotary_interleaved,
             mscale=mscale,
         ).squeeze(1)
-    else:
-        # CASE 2: Traditional mapping without offsets
-        # Build packed freqs for all sequences using the standard mapping, then apply once
-        sequence_splits = torch.split(t, seqlens)
-        freqs_packed = torch.cat(
-            [_get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs) for x in sequence_splits],
-            dim=0,
-        )
 
-        return _apply_rotary_pos_emb_bshd(
-            t.unsqueeze(1),
-            freqs_packed,
+    # CASE 2: Traditional mapping without offsets
+    output = torch.empty_like(t)
+    output_offset = 0
+    for x in sequence_splits:
+        freq_slice = _get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs)
+        output_slice = _apply_rotary_pos_emb_bshd(
+            x.unsqueeze(1),
+            freq_slice,
             rotary_interleaved=rotary_interleaved,
             mla_rotary_interleaved=mla_rotary_interleaved,
             mscale=mscale,
         ).squeeze(1)
+        output.narrow(0, output_offset, x.size(0)).copy_(output_slice)
+        output_offset += x.size(0)
+
+    return output
 
 
 def apply_rotary_pos_emb(
@@ -283,6 +283,8 @@ def apply_rotary_pos_emb(
     # Keep for backward compatibility. Will deprecate in the future.
     if cp_group is None:
         cp_group = parallel_state.get_context_parallel_group()
+    if mla_rotary_interleaved is None:
+        mla_rotary_interleaved = config.multi_latent_attention
 
     if config.apply_rope_fusion:
         if cu_seqlens is None:
diff --git a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py
index 8231a2a3764..3f14aeaa9df 100644
--- a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py
+++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py
@@ -6,6 +6,10 @@
 from megatron.core.models.backends import BackendSpecProvider
 from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNetSubmodules
 from megatron.core.transformer.enums import AttnMaskType, LayerType
+from megatron.core.transformer.experimental_attention_variant.absorbed_mla import (
+    AbsorbedMLASelfAttention,
+    AbsorbedMLASelfAttentionSubmodules,
+)
 from megatron.core.transformer.experimental_attention_variant.dsa import (
     DSAIndexer,
     DSAIndexerSubmodules,
@@ -13,10 +17,6 @@
     DSAttentionSubmodules,
 )
 from megatron.core.transformer.identity_op import IdentityOp
-from megatron.core.transformer.multi_latent_attention import (
-    MLASelfAttention,
-    MLASelfAttentionSubmodules,
-)
 from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.transformer_block import (
     TransformerBlockSubmodules,
@@ -109,9 +109,9 @@ def get_dsa_module_spec_for_backend(
     )
 
     attention = ModuleSpec(
-        module=MLASelfAttention,
+        module=AbsorbedMLASelfAttention,
         params={"attn_mask_type": AttnMaskType.causal},
-        submodules=MLASelfAttentionSubmodules(
+        submodules=AbsorbedMLASelfAttentionSubmodules(
             linear_q_proj=backend.column_parallel_linear(),
             linear_q_down_proj=backend.linear(),
             linear_q_up_proj=backend.column_parallel_linear(),
diff --git a/megatron/core/models/hybrid/hybrid_layer_specs.py b/megatron/core/models/hybrid/hybrid_layer_specs.py
index 5b968f720c0..e1624293b5a 100755
--- a/megatron/core/models/hybrid/hybrid_layer_specs.py
+++ b/megatron/core/models/hybrid/hybrid_layer_specs.py
@@ -26,6 +26,10 @@
 )
 from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
 from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.experimental_attention_variant.absorbed_mla import (
+    AbsorbedMLASelfAttention,
+    AbsorbedMLASelfAttentionSubmodules,
+)
 from megatron.core.transformer.experimental_attention_variant.dsa import (
     DSAIndexer,
     DSAIndexerSubmodules,
@@ -135,9 +139,9 @@
             submodules=TransformerLayerSubmodules(
                 input_layernorm=TENorm,
                 self_attention=ModuleSpec(
-                    module=MLASelfAttention,
+                    module=AbsorbedMLASelfAttention,
                     params={"attn_mask_type": AttnMaskType.causal},
-                    submodules=MLASelfAttentionSubmodules(
+                    submodules=AbsorbedMLASelfAttentionSubmodules(
                         linear_q_proj=TEColumnParallelLinear,
                         linear_q_down_proj=TELinear,
                         linear_q_up_proj=TEColumnParallelLinear,
diff --git a/megatron/core/pipeline_parallel/hybrid_cp_schedule.py b/megatron/core/pipeline_parallel/hybrid_cp_schedule.py
index 27b5fc87945..96b2ac6b553 100644
--- a/megatron/core/pipeline_parallel/hybrid_cp_schedule.py
+++ b/megatron/core/pipeline_parallel/hybrid_cp_schedule.py
@@ -545,9 +545,15 @@ def _get_new_data_iterator(sample_id_in_group, group_id):
             )
             sample["local_cp_size"] = torch.tensor(partner_cp_size, dtype=torch.int32)
             new_data_iterator = RerunDataIterator(iter([sample]))
-            return new_data_iterator
         else:
-            return None
+            partner_cp_size = 0
+            new_data_iterator = None
+
+        partner_cp_size_tensor = torch.tensor(
+            [partner_cp_size], dtype=torch.int32, device=torch.cuda.current_device()
+        )
+        _broadcast(partner_cp_size_tensor)
+        return new_data_iterator, int(partner_cp_size_tensor.item())
 
     # We get data once per global batch and schedule the sub-samples.
     # TODO(pmannan): Should we wrap the data_iterator here instead of the training.py file?
@@ -579,7 +585,7 @@ def _get_new_data_iterator(sample_id_in_group, group_id):
             sample_ids_this_group = sample_id_groups[j][hdp_rank] if is_first_tp_rank else None
             for i in range(num_samples_this_group[j]):
                 # Call forward step for each sub-sample
-                new_data_iterator = _get_new_data_iterator(i, j)
+                new_data_iterator, cp_group_size = _get_new_data_iterator(i, j)
                 # TODO: Find the usage of current_microbatch and is_first_microbatch and
                 # how that may affect my usage.
                 output_tensor, num_tokens = forward_step(
@@ -590,7 +596,8 @@ def _get_new_data_iterator(sample_id_in_group, group_id):
                     input_tensor,
                     forward_data_store,
                     config,
-                    collect_non_loss_data,
+                    cp_group_size=cp_group_size,
+                    collect_non_loss_data=collect_non_loss_data,
                     is_first_microbatch=check_first_val_step(
                         first_val_step, forward_only, current_microbatch == 0
                     ),
@@ -599,9 +606,7 @@ def _get_new_data_iterator(sample_id_in_group, group_id):
                 current_microbatch += 1
                 total_num_tokens += num_tokens.item()
                 if not forward_only:
-                    backward_step(
-                        input_tensor, output_tensor, output_tensor_grad, model_type, config
-                    )
+                    backward_step(input_tensor, output_tensor, output_tensor_grad, config)
 
             # Create a barrier at end of each group.
             # This barrier ensures that all ranks are prepared to change assigned CP group sizes and
@@ -614,7 +619,7 @@ def _get_new_data_iterator(sample_id_in_group, group_id):
     with no_sync_func():
         sample_ids_this_group = sample_id_groups[-1][hdp_rank] if is_first_tp_rank else None
         for i in range(num_samples_this_group[-1] - 1):
-            new_data_iterator = _get_new_data_iterator(i, -1)
+            new_data_iterator, cp_group_size = _get_new_data_iterator(i, -1)
             # Call forward step for each sub-sample
             output_tensor, num_tokens = forward_step(
                 forward_step_func,
@@ -624,7 +629,8 @@ def _get_new_data_iterator(sample_id_in_group, group_id):
                 input_tensor,
                 forward_data_store,
                 config,
-                collect_non_loss_data,
+                cp_group_size=cp_group_size,
+                collect_non_loss_data=collect_non_loss_data,
                 is_first_microbatch=check_first_val_step(
                     first_val_step, forward_only, current_microbatch == 0
                 ),
@@ -633,11 +639,11 @@ def _get_new_data_iterator(sample_id_in_group, group_id):
             current_microbatch += 1
             total_num_tokens += num_tokens.item()
             if not forward_only:
-                backward_step(input_tensor, output_tensor, output_tensor_grad, model_type, config)
+                backward_step(input_tensor, output_tensor, output_tensor_grad, config)
 
     # The last sub-sample of the last group of the last microbatch is
     # run out of the context handler.
-    new_data_iterator = _get_new_data_iterator(-1, -1)
+    new_data_iterator, cp_group_size = _get_new_data_iterator(-1, -1)
     # Call forward step for each sub-sample
     output_tensor, num_tokens = forward_step(
         forward_step_func,
@@ -647,7 +653,8 @@ def _get_new_data_iterator(sample_id_in_group, group_id):
         input_tensor,
         forward_data_store,
         config,
-        collect_non_loss_data,
+        cp_group_size=cp_group_size,
+        collect_non_loss_data=collect_non_loss_data,
         is_first_microbatch=check_first_val_step(
             first_val_step, forward_only, current_microbatch == 0
         ),
@@ -655,6 +662,6 @@ def _get_new_data_iterator(sample_id_in_group, group_id):
     )
     total_num_tokens += num_tokens.item()
     if not forward_only:
-        backward_step(input_tensor, output_tensor, output_tensor_grad, model_type, config)
+        backward_step(input_tensor, output_tensor, output_tensor_grad, config)
 
     return forward_data_store, total_num_tokens
diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py
index b2c23807bea..d13117109d5 100644
--- a/megatron/core/pipeline_parallel/schedules.py
+++ b/megatron/core/pipeline_parallel/schedules.py
@@ -226,28 +226,58 @@ def get_tensor_device(tensor: Union[torch.Tensor, Dict[str, torch.Tensor]]):
     return tensor.device
 
 
-def _get_mtp_loss_scale(config, device: torch.device) -> torch.Tensor:
-    """Get the MTP loss scale on the output tensor device."""
+def _normalize_loss_scale(loss_scale, device: torch.device, scale_func_name: str) -> torch.Tensor:
+    """Normalize loss scale outputs to a size-1 tensor on the output tensor device."""
+    loss_scale = torch.as_tensor(loss_scale, device=device)
+    if loss_scale.numel() != 1:
+        raise ValueError(
+            f"{scale_func_name} must return a scalar or size-1 tensor for loss scaling, "
+            f"but returned a tensor with {loss_scale.numel()} elements."
+        )
+    return loss_scale
 
-    def _normalize_loss_scale(loss_scale, scale_func_name: str) -> torch.Tensor:
-        loss_scale = torch.as_tensor(loss_scale, device=device)
-        if loss_scale.numel() != 1:
-            raise ValueError(
-                f"{scale_func_name} must return a scalar or size-1 tensor for MTP loss scaling, "
-                f"but returned a tensor with {loss_scale.numel()} elements."
-            )
-        return loss_scale
 
-    mtp_grad_scale_func = getattr(config, 'mtp_grad_scale_func', None)
-    if mtp_grad_scale_func is not None:
-        return _normalize_loss_scale(mtp_grad_scale_func(), "mtp_grad_scale_func")
+def _compute_loss_scale(config, device: torch.device) -> torch.Tensor:
+    """Calculate the loss scale from grad_scale_func or default to 1."""
     if config.grad_scale_func is not None:
         return _normalize_loss_scale(
-            config.grad_scale_func(torch.ones(1, device=device)), "grad_scale_func"
+            config.grad_scale_func(torch.ones(1, device=device)), device, "grad_scale_func"
         )
     return torch.ones(1, device=device)
 
 
+def _get_moe_loss_scale(config, device: torch.device) -> torch.Tensor:
+    """Get the MoE loss scale on the output tensor device."""
+    moe_grad_scale_func = getattr(config, 'moe_grad_scale_func', None)
+    if moe_grad_scale_func is not None:
+        return _normalize_loss_scale(moe_grad_scale_func(), device, "moe_grad_scale_func")
+    return _compute_loss_scale(config, device)
+
+
+def _get_mtp_loss_scale(config, device: torch.device) -> torch.Tensor:
+    """Get the MTP loss scale on the output tensor device."""
+    mtp_grad_scale_func = getattr(config, 'mtp_grad_scale_func', None)
+    if mtp_grad_scale_func is not None:
+        return _normalize_loss_scale(mtp_grad_scale_func(), device, "mtp_grad_scale_func")
+    return _compute_loss_scale(config, device)
+
+
+def _get_experimental_attention_variant_loss_scale_func(config):
+    """Get the loss scale hook for experimental attention variants."""
+    loss_scale_func = getattr(config, 'experimental_attention_variant_loss_scale_func', None)
+    if loss_scale_func is not None:
+        return loss_scale_func
+
+    if getattr(config, 'experimental_attention_variant', None) == 'dsa':
+        from megatron.core.transformer.experimental_attention_variant.dsa import (
+            DSAIndexerLossAutoScaler,
+        )
+
+        return DSAIndexerLossAutoScaler.set_loss_scale
+
+    return None
+
+
 def forward_step_calc_loss(
     model,
     output_tensor,
@@ -262,9 +292,6 @@ def forward_step_calc_loss(
 ):
     """Calculate the loss and number of tokens for forward_step()"""
 
-    from megatron.core.transformer.experimental_attention_variant.dsa import (
-        DSAIndexerLossAutoScaler,
-    )
     from megatron.core.transformer.multi_token_prediction import MTPLossAutoScaler
 
     model_vp_stage = getattr(model, "vp_stage", None)
@@ -315,16 +342,8 @@ def forward_step_calc_loss(
     # Since we use a trick to do backward on the auxiliary loss, we need to set the scale
     # explicitly.
     if hasattr(config, 'num_moe_experts') and config.num_moe_experts is not None:
-        # Calculate the loss scale based on moe_grad_scale_func (preferred),
-        # grad_scale_func (fallback), or default to 1.
         device = get_tensor_device(output_tensor)
-        moe_grad_scale_func = getattr(config, 'moe_grad_scale_func', None)
-        if moe_grad_scale_func is not None:
-            loss_scale = moe_grad_scale_func()
-        elif config.grad_scale_func is not None:
-            loss_scale = config.grad_scale_func(torch.ones(1, device=device))
-        else:
-            loss_scale = torch.ones(1, device=device)
+        loss_scale = _get_moe_loss_scale(config, device)
         # Set the loss scale
         if config.calculate_per_token_loss:
             MoEAuxLossAutoScaler.set_loss_scale(loss_scale)
@@ -344,17 +363,20 @@ def forward_step_calc_loss(
         else:
             MTPLossAutoScaler.set_loss_scale(loss_scale / num_microbatches)
 
-    # Set the loss scale for DSA (Dynamic Sparse Attention) indexer loss.
-    if getattr(config, 'experimental_attention_variant', None) == 'dsa':
-        loss_scale = (
-            config.grad_scale_func(torch.ones(1, device=output_tensor.device))
-            if config.grad_scale_func is not None
-            else torch.ones(1, device=output_tensor.device)
-        )
+    # Set the loss scale for any experimental attention-variant auxiliary loss.
+    experimental_attention_variant_loss_scale_func = (
+        _get_experimental_attention_variant_loss_scale_func(config)
+    )
+    if experimental_attention_variant_loss_scale_func is not None:
+        device = get_tensor_device(output_tensor)
+        loss_scale = _compute_loss_scale(config, device)
         if config.calculate_per_token_loss:
-            DSAIndexerLossAutoScaler.set_loss_scale(loss_scale)
+            experimental_attention_variant_loss_scale_func(loss_scale)
         else:
-            DSAIndexerLossAutoScaler.set_loss_scale(loss_scale / num_microbatches)
+            cp_size_for_scaling = cp_group_size if cp_group_size is not None else 1
+            experimental_attention_variant_loss_scale_func(
+                loss_scale * cp_size_for_scaling / num_microbatches
+            )
 
     return output_tensor, num_tokens
 
diff --git a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py
index 48e6c76ea2f..f8f4c37ea7c 100644
--- a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py
+++ b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py
@@ -18,6 +18,7 @@
 
 import torch
 
+from megatron.core import tensor_parallel
 from megatron.core.extensions.transformer_engine import HAVE_TE
 from megatron.core.models.common.embeddings import (
     RotaryEmbedding,
@@ -36,7 +37,7 @@
 from megatron.core.transformer.enums import AttnMaskType
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
 from megatron.core.transformer.transformer_config import MLATransformerConfig
-from megatron.core.utils import deprecate_inference_params, get_pg_size
+from megatron.core.utils import deprecate_inference_params, get_pg_size, is_te_min_version
 
 try:
     from megatron.core.fusions.fused_mla_yarn_rope_apply import (
@@ -58,6 +59,52 @@
     TEColumnParallelLinear, TELinear, Linear, set_save_original_input = None, None, None, None
 
 
+def _restore_packed_thd_batch_dim(
+    core_attn_out: torch.Tensor, hidden_states: torch.Tensor, packed_seq_params
+) -> torch.Tensor:
+    """Restore the singleton packed-THD batch dim only when core attention omitted it."""
+    if (
+        packed_seq_params is not None
+        and packed_seq_params.qkv_format == 'thd'
+        and core_attn_out.ndim == hidden_states.ndim - 1
+    ):
+        core_attn_out = core_attn_out.unsqueeze(1)
+    return core_attn_out
+
+
+def _apply_absorbed_v_up_projection(
+    core_attn_out: torch.Tensor,
+    v_up_weight: torch.Tensor,
+    num_attention_heads_per_partition: int,
+    kv_lora_rank: int,
+    v_head_dim: int,
+    core_consumed_v_up_projection: bool,
+) -> torch.Tensor:
+    """Apply V up projection unless core attention already consumed the projection weight."""
+    latent_output_size = num_attention_heads_per_partition * kv_lora_rank
+    projected_output_size = num_attention_heads_per_partition * v_head_dim
+    if core_consumed_v_up_projection:
+        if core_attn_out.size(-1) != projected_output_size:
+            raise RuntimeError(
+                "AbsorbedMLA core attention returned unexpected projected hidden size: "
+                f"{core_attn_out.size(-1)}. Expected projected={projected_output_size}."
+            )
+        return core_attn_out
+
+    if core_attn_out.size(-1) != latent_output_size:
+        raise RuntimeError(
+            "AbsorbedMLA core attention returned unexpected hidden size: "
+            f"{core_attn_out.size(-1)}. Expected latent={latent_output_size}."
+        )
+
+    core_attn_out = core_attn_out.view(
+        *core_attn_out.shape[:-1], num_attention_heads_per_partition, kv_lora_rank
+    )
+    core_attn_out = torch.einsum("...nc,ndc->...nd", core_attn_out, v_up_weight)
+    core_attn_out = core_attn_out.contiguous()
+    return core_attn_out.view(*core_attn_out.shape[:-2], -1)
+
+
 @dataclass
 class AbsorbedMLASelfAttentionSubmodules:
     """
@@ -68,8 +115,7 @@ class AbsorbedMLASelfAttentionSubmodules:
     linear_q_down_proj: Union[ModuleSpec, type] = None
     linear_q_up_proj: Union[ModuleSpec, type] = None
     linear_kv_down_proj: Union[ModuleSpec, type] = None
-    linear_k_up_proj: Union[ModuleSpec, type] = None
-    linear_v_up_proj: Union[ModuleSpec, type] = None
+    linear_kv_up_proj: Union[ModuleSpec, type] = None
     core_attention: Union[ModuleSpec, type] = None
     linear_proj: Union[ModuleSpec, type] = None
     q_layernorm: Union[ModuleSpec, type] = None
@@ -110,12 +156,16 @@ def __init__(
             layer_number=layer_number,
             attn_mask_type=attn_mask_type,
             attention_type="self",
+            cp_comm_type=cp_comm_type,
             pg_collection=pg_collection,
             pp_layer_offset=pp_layer_offset,
             name=name,
         )
 
         assert not config.add_bias_linear, "add_bias_linear is not supported for AbsorbedMLA"
+        assert not (
+            config.tensor_model_parallel_size > 1 and not config.sequence_parallel
+        ), "AbsorbedMLA requires sequence_parallel when tensor_model_parallel_size > 1"
 
         self.query_projection_size = self.config.v_head_dim * self.config.num_attention_heads
         self.q_head_dim = self.config.qk_head_dim + self.config.qk_pos_emb_head_dim
@@ -302,34 +352,19 @@ def __init__(
             **kv_down_proj_kwargs,
         )
 
-        # Build separate K and V up projections
-        self.linear_k_up_proj = build_module(
-            submodules.linear_k_up_proj,
+        self.linear_kv_up_proj = build_module(
+            submodules.linear_kv_up_proj,
             self.config.kv_lora_rank,
-            self.config.num_attention_heads * self.config.qk_head_dim,
+            self.config.num_attention_heads * (self.config.qk_head_dim + self.config.v_head_dim),
             config=self.config,
             init_method=self.config.init_method,
             gather_output=False,
             bias=False,
             skip_bias_add=False,
             is_expert=False,
-            tp_comm_buffer_name='k_up_proj',
+            tp_comm_buffer_name='kv_up_proj',
             tp_group=pg_collection.tp,
-            name=(name + ".linear_k_up_proj") if name is not None else None,
-        )
-        self.linear_v_up_proj = build_module(
-            submodules.linear_v_up_proj,
-            self.config.kv_lora_rank,
-            self.config.num_attention_heads * self.config.v_head_dim,
-            config=self.config,
-            init_method=self.config.init_method,
-            gather_output=False,
-            bias=False,
-            skip_bias_add=False,
-            is_expert=False,
-            tp_comm_buffer_name='v_up_proj',
-            tp_group=pg_collection.tp,
-            name=(name + ".linear_v_up_proj") if name is not None else None,
+            name=(name + ".linear_kv_up_proj") if name is not None else None,
         )
 
         if self.config.q_lora_rank is not None:
@@ -511,17 +546,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po
             # [num_tokens, qk_pos_emb_head_dim] -> [num_tokens, 1, qk_pos_emb_head_dim]
             k_pos_emb = torch.unsqueeze(k_pos_emb, -2)
 
-            # Prepare k_up_weight for absorption
-            # k_up_weight: linear_k_up_proj.weight viewed as [n, qk_head_dim, kv_lora_rank]
-            assert self.linear_k_up_proj.weight.size(0) == (
-                self.num_attention_heads_per_partition * self.config.qk_head_dim
-            )
-            assert self.linear_k_up_proj.weight.size(1) == self.config.kv_lora_rank
-            k_up_weight = self.linear_k_up_proj.weight.view(
-                self.num_attention_heads_per_partition,
-                self.config.qk_head_dim,
-                self.config.kv_lora_rank,
-            )
+            k_up_weight, _ = self._get_kv_up_weights()
 
             if self.config.apply_rope_fusion:
                 # q_no_pe: [num_tokens, n, qk_head_dim]
@@ -643,17 +668,77 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po
 
         return q_absorbed, kv_compressed, q_compressed
 
+    def _get_v_up_weight(self) -> torch.Tensor:
+        """Return V up-projection weight in per-head layout."""
+        _, v_up_weight = self._get_kv_up_weights()
+        return v_up_weight
+
+    def _get_kv_up_weights(self) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return K and V up-projection weights from the combined per-head MLA layout."""
+        expected_rows = self.num_attention_heads_per_partition * (
+            self.config.qk_head_dim + self.config.v_head_dim
+        )
+        assert self.linear_kv_up_proj.weight.size(0) == expected_rows
+        assert self.linear_kv_up_proj.weight.size(1) == self.config.kv_lora_rank
+        kv_up_weight = self.linear_kv_up_proj.weight.view(
+            self.num_attention_heads_per_partition,
+            self.config.qk_head_dim + self.config.v_head_dim,
+            self.config.kv_lora_rank,
+        )
+        k_up_weight = kv_up_weight[:, : self.config.qk_head_dim, :]
+        v_up_weight = kv_up_weight[:, self.config.qk_head_dim :, :]
+        return k_up_weight, v_up_weight
+
+    def _combine_split_kv_up_weights(
+        self, k_up_weight: torch.Tensor, v_up_weight: torch.Tensor
+    ) -> torch.Tensor:
+        """Combine pre-refactor split K/V up-projection weights into the new layout."""
+        num_heads = self.num_attention_heads_per_partition
+        qk_head_dim = self.config.qk_head_dim
+        v_head_dim = self.config.v_head_dim
+        kv_lora_rank = self.config.kv_lora_rank
+
+        k_up_weight = k_up_weight.view(num_heads, qk_head_dim, kv_lora_rank)
+        v_up_weight = v_up_weight.view(num_heads, v_head_dim, kv_lora_rank)
+        return (
+            torch.cat((k_up_weight, v_up_weight), dim=1)
+            .contiguous()
+            .view(num_heads * (qk_head_dim + v_head_dim), kv_lora_rank)
+        )
+
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        """Load checkpoints saved with either combined or split K/V up-projection weights."""
+        combined_key = f"{prefix}linear_kv_up_proj.weight"
+        k_up_key = f"{prefix}linear_k_up_proj.weight"
+        v_up_key = f"{prefix}linear_v_up_proj.weight"
+        if combined_key not in state_dict and k_up_key in state_dict and v_up_key in state_dict:
+            state_dict[combined_key] = self._combine_split_kv_up_weights(
+                state_dict.pop(k_up_key), state_dict.pop(v_up_key)
+            )
+
+        combined_extra_state_key = f"{prefix}linear_kv_up_proj._extra_state"
+        k_up_extra_state_key = f"{prefix}linear_k_up_proj._extra_state"
+        v_up_extra_state_key = f"{prefix}linear_v_up_proj._extra_state"
+        if k_up_extra_state_key in state_dict or v_up_extra_state_key in state_dict:
+            k_extra_state = state_dict.pop(k_up_extra_state_key, None)
+            v_extra_state = state_dict.pop(v_up_extra_state_key, None)
+            if combined_extra_state_key not in state_dict:
+                state_dict[combined_extra_state_key] = (
+                    k_extra_state if k_extra_state is not None else v_extra_state
+                )
+
+        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
     def _checkpointed_attention_forward(
         self,
         q_absorbed,
         k_compressed,
-        v_compressed,
         hidden_states,
         q_compressed,
         attention_mask,
-        rotary_pos_emb=None,
+        up_v_weight,
+        position_ids=None,
         attn_mask_type=None,
-        attention_bias=None,
         packed_seq_params=None,
     ):
         """Forward method with selective activation checkpointing."""
@@ -661,23 +746,22 @@ def _checkpointed_attention_forward(
         def custom_forward(*inputs):
             q_absorbed = inputs[0]
             k_compressed = inputs[1]
-            v_compressed = inputs[2]
-            hidden_states = inputs[3]
-            q_compressed = inputs[4]
-            attention_mask = inputs[5]
-            attn_mask_type = inputs[7]
-            attention_bias = inputs[8]
-            packed_seq_params = inputs[9]
+            hidden_states = inputs[2]
+            q_compressed = inputs[3]
+            attention_mask = inputs[4]
+            up_v_weight = inputs[5]
+            attn_mask_type = inputs[6]
             attn_mask_type = AttnMaskType(attn_mask_type.item())
             output_ = self.core_attention(
                 q_absorbed,
                 k_compressed,
-                v_compressed,
-                hidden_states,
-                q_compressed,
+                None,
                 attention_mask,
+                x=hidden_states,
+                qr=q_compressed,
+                up_v_weight=up_v_weight,
+                position_ids=position_ids,
                 attn_mask_type=attn_mask_type,
-                attention_bias=attention_bias,
                 packed_seq_params=packed_seq_params,
             )
             return output_
@@ -690,14 +774,11 @@ def custom_forward(*inputs):
             False,
             q_absorbed,
             k_compressed,
-            v_compressed,
             hidden_states,
             q_compressed,
             attention_mask,
-            rotary_pos_emb,
+            up_v_weight,
             attn_mask_type,
-            attention_bias,
-            packed_seq_params,
         )
 
         return hidden_states
@@ -714,6 +795,7 @@ def forward(
         rotary_pos_cos_sin=None,
         attention_bias=None,
         packed_seq_params=None,
+        position_ids=None,
         sequence_len_offset=None,
         *,
         inference_params=None,
@@ -742,6 +824,7 @@ def forward(
         assert q_absorbed.is_contiguous()
         assert q_compressed.is_contiguous()
         assert kv_compressed.is_contiguous()
+        v_up_weight = self._get_v_up_weight()
 
         # ==================================
         # Core attention computation
@@ -750,10 +833,11 @@ def forward(
             core_attn_out = self._checkpointed_attention_forward(
                 q_absorbed,
                 kv_compressed,
-                None,
                 hidden_states,
                 q_compressed,
                 attention_mask,
+                v_up_weight,
+                position_ids=position_ids,
                 packed_seq_params=packed_seq_params,
             )
         else:
@@ -761,9 +845,11 @@ def forward(
                 q_absorbed,
                 kv_compressed,
                 None,
-                hidden_states,
-                q_compressed,
                 attention_mask,
+                x=hidden_states,
+                qr=q_compressed,
+                up_v_weight=v_up_weight,
+                position_ids=position_ids,
                 packed_seq_params=packed_seq_params,
                 attn_mask_type=self.attn_mask_type,
             )
@@ -771,24 +857,21 @@ def forward(
         # ==================================
         # Apply V up projection
         # ==================================
-        assert self.linear_v_up_proj.weight.size(0) == (
-            self.num_attention_heads_per_partition * self.config.v_head_dim
-        )
-        assert self.linear_v_up_proj.weight.size(1) == self.config.kv_lora_rank
-        v_up_weight = self.linear_v_up_proj.weight.view(
-            self.num_attention_heads_per_partition, self.config.v_head_dim, self.config.kv_lora_rank
+        core_consumed_v_up_projection = getattr(
+            self.core_attention, "consumes_absorbed_v_up_projection", False
         )
-        core_attn_out = core_attn_out.view(
-            *core_attn_out.shape[:-1],
+        core_attn_out = _apply_absorbed_v_up_projection(
+            core_attn_out,
+            v_up_weight,
             self.num_attention_heads_per_partition,
             self.config.kv_lora_rank,
+            self.config.v_head_dim,
+            core_consumed_v_up_projection,
         )
-        core_attn_out = torch.einsum("...nc,ndc->...nd", core_attn_out, v_up_weight)
-        core_attn_out = core_attn_out.contiguous()
-        core_attn_out = core_attn_out.view(*core_attn_out.shape[:-2], -1)
 
-        if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd':
-            core_attn_out = core_attn_out.unsqueeze(1)
+        core_attn_out = _restore_packed_thd_batch_dim(
+            core_attn_out, hidden_states, packed_seq_params
+        )
 
         assert core_attn_out.ndim == hidden_states.ndim
         assert core_attn_out.shape[0] == (
@@ -823,8 +906,7 @@ def backward_dw(self) -> NoReturn:
 
     def _backward_kv_proj(self):
         """Computes weight gradients of KV projection layers."""
-        self.linear_k_up_proj.backward_dw()
-        self.linear_v_up_proj.backward_dw()
+        self.linear_kv_up_proj.backward_dw()
         self.linear_kv_down_proj.backward_dw()
 
     def _backward_q_proj(self):
@@ -854,115 +936,3 @@ def clip_qk(self):
         function after Muon optimizer step.
         """
         raise NotImplementedError("clip_qk is not implemented for AbsorbedMLA")
-
-    def _combine_kv_weights(self, k_weight, v_weight):
-        """Combine separate K and V weights into MLA's interleaved format.
-
-        MLA's linear_kv_up_proj weight layout (per head interleaved):
-            [head0_K, head0_V, head1_K, head1_V, ...]
-
-        AbsorbedMLA's separate weights layout:
-            K: [head0_K, head1_K, ...]
-            V: [head0_V, head1_V, ...]
-
-        This method interleaves K and V per head to match MLA's format.
-
-        Args:
-            k_weight: [num_heads_per_partition * qk_head_dim, kv_lora_rank]
-            v_weight: [num_heads_per_partition * v_head_dim, kv_lora_rank]
-
-        Returns:
-            combined: [num_heads_per_partition * (qk_head_dim + v_head_dim), kv_lora_rank]
-        """
-        n = self.num_attention_heads_per_partition
-        qk_dim = self.config.qk_head_dim
-        v_dim = self.config.v_head_dim
-        lora_rank = self.config.kv_lora_rank
-
-        # Reshape to per-head format
-        k_per_head = k_weight.view(n, qk_dim, lora_rank)
-        v_per_head = v_weight.view(n, v_dim, lora_rank)
-
-        # Concatenate K and V for each head along dim=1
-        # Result: [n, qk_dim + v_dim, lora_rank]
-        combined_per_head = torch.cat([k_per_head, v_per_head], dim=1)
-
-        # Reshape back to linear weight format
-        combined_weight = combined_per_head.view(n * (qk_dim + v_dim), lora_rank)
-
-        return combined_weight
-
-    def _split_kv_weights(self, combined_weight):
-        """Split MLA's interleaved KV weight into separate K and V weights.
-
-        MLA's linear_kv_up_proj weight layout (per head interleaved):
-            [head0_K, head0_V, head1_K, head1_V, ...]
-
-        This method extracts K and V into separate tensors:
-            K: [head0_K, head1_K, ...]
-            V: [head0_V, head1_V, ...]
-
-        Args:
-            combined_weight: [num_heads_per_partition * (qk_head_dim + v_head_dim), kv_lora_rank]
-
-        Returns:
-            k_weight: [num_heads_per_partition * qk_head_dim, kv_lora_rank]
-            v_weight: [num_heads_per_partition * v_head_dim, kv_lora_rank]
-        """
-        n = self.num_attention_heads_per_partition
-        qk_dim = self.config.qk_head_dim
-        v_dim = self.config.v_head_dim
-        lora_rank = self.config.kv_lora_rank
-
-        # Reshape to per-head format
-        combined_per_head = combined_weight.view(n, qk_dim + v_dim, lora_rank)
-
-        # Split K and V for each head (slicing creates non-contiguous views)
-        k_per_head = combined_per_head[:, :qk_dim, :]  # [n, qk_dim, lora_rank]
-        v_per_head = combined_per_head[:, qk_dim:, :]  # [n, v_dim, lora_rank]
-
-        # Make contiguous and reshape back to linear weight format
-        k_weight = k_per_head.contiguous().view(n * qk_dim, lora_rank)
-        v_weight = v_per_head.contiguous().view(n * v_dim, lora_rank)
-
-        return k_weight, v_weight
-
-    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
-        """Handle loading from checkpoints with combined KV up projection weights.
-
-        This method splits the combined 'linear_kv_up_proj.weight' (which has per-head
-        interleaved K and V) into separate 'linear_k_up_proj.weight' and 'linear_v_up_proj.weight'.
-        """
-        combined_key = f'{prefix}linear_kv_up_proj.weight'
-        k_up_key = f'{prefix}linear_k_up_proj.weight'
-        v_up_key = f'{prefix}linear_v_up_proj.weight'
-
-        # Split combined KV weights into separate K and V
-        if combined_key in state_dict:
-            combined_weight = state_dict[combined_key]
-
-            # Split with proper per-head de-interleaving
-            k_weight, v_weight = self._split_kv_weights(combined_weight)
-
-            state_dict[k_up_key] = k_weight
-            state_dict[v_up_key] = v_weight
-
-            del state_dict[combined_key]
-
-        combined_extra_state_key = f'{prefix}linear_kv_up_proj._extra_state'
-        k_up_extra_state_key = f'{prefix}linear_k_up_proj._extra_state'
-        v_up_extra_state_key = f'{prefix}linear_v_up_proj._extra_state'
-
-        if combined_extra_state_key in state_dict:
-            combined_extra_state = state_dict[combined_extra_state_key]
-
-            assert isinstance(combined_extra_state, torch.Tensor)
-            # Now we can only handle the case where the extra state is empty.
-            assert combined_extra_state.numel() == 0
-
-            state_dict[k_up_extra_state_key] = combined_extra_state.clone()
-            state_dict[v_up_extra_state_key] = combined_extra_state.clone()
-
-            del state_dict[combined_extra_state_key]
-
-        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py
index 5c5f77363dc..24cd5ba86a6 100644
--- a/megatron/core/transformer/experimental_attention_variant/dsa.py
+++ b/megatron/core/transformer/experimental_attention_variant/dsa.py
@@ -17,6 +17,11 @@
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region
 from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.experimental_attention_variant import (
+    dsa_kernels,
+    dsa_layout,
+    dsa_masking,
+)
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
 from megatron.core.transformer.transformer_config import TransformerConfig
@@ -27,6 +32,187 @@
     hadamard_transform = None
 
 
+def _unfused_absorbed_dsa_fn(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    topk_indices: torch.Tensor,
+    softmax_scale: float,
+    v_channels: int,
+    mask: Optional[torch.Tensor] = None,
+    varlen_starts: Optional[torch.Tensor] = None,
+    varlen_ends: Optional[torch.Tensor] = None,
+    key_positions: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Unfused absorbed-MLA attention: output stays [sq, b, np, v_channels]."""
+    sq, b, np, hn = query.size()
+    skv = key.size(0)
+    assert key.size(2) == 1, "Absorbed DSA expects MQA key head dimension = 1"
+    assert key.size(-1) >= v_channels, "key last dim must contain latent value channels"
+    row_mask, varlen_starts, varlen_ends, key_positions = dsa_masking.prepare_sparse_mask_context(
+        mask=mask,
+        varlen_starts=varlen_starts,
+        varlen_ends=varlen_ends,
+        key_positions=key_positions,
+        sq=sq,
+        sk=skv,
+        b=b,
+        device=query.device,
+    )
+
+    # [sq,b,np,hn] -> [b,np,sq,hn]
+    q = query.permute(1, 2, 0, 3)
+    # [skv,b,1,hn] -> [b,1,hn,skv]
+    k = key.permute(1, 2, 3, 0)
+    attention_scores = torch.matmul(q.float(), k.float()) * softmax_scale
+
+    # Sparse + causal/varlen validity mask.
+    index_mask = torch.full((b, sq, skv), float("-inf"), device=attention_scores.device)
+    dsa_masking.scatter_topk_into_index_mask(index_mask, topk_indices, seq_chunk_size=256)
+    index_mask = dsa_masking.apply_sparse_validity_to_index_mask(
+        index_mask,
+        row_mask=row_mask,
+        varlen_starts=varlen_starts,
+        varlen_ends=varlen_ends,
+        key_positions=key_positions,
+    )
+
+    attention_scores += index_mask.unsqueeze(1)
+    attention_scores = torch.nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32)
+
+    # Latent value is the first v_channels slice of absorbed key cache.
+    value = key[..., :v_channels].permute(1, 2, 0, 3)  # [b,1,skv,v]
+    output = torch.matmul(attention_scores.to(value.dtype), value)  # [b,np,sq,v]
+    return output.permute(2, 0, 1, 3).contiguous()
+
+
+def _run_sparse_attention(
+    *,
+    absorbed_mla: bool,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: Optional[torch.Tensor],
+    up_v_weight: Optional[torch.Tensor],
+    topk_indices: torch.Tensor,
+    softmax_scale: float,
+    config: TransformerConfig,
+    mask: Optional[torch.Tensor],
+    varlen_starts: Optional[torch.Tensor],
+    varlen_ends: Optional[torch.Tensor],
+    key_positions: Optional[torch.Tensor],
+) -> torch.Tensor:
+    """Run sparse attention for absorbed and non-absorbed MLA paths."""
+    if absorbed_mla:
+        latent_v_channels = int(getattr(config, "kv_lora_rank", 0) or 0)
+        if latent_v_channels <= 0:
+            raise RuntimeError(
+                "Invalid kv_lora_rank for absorbed-MLA DSAttention sparse attention."
+            )
+        if up_v_weight is None:
+            raise RuntimeError(
+                "Absorbed DSAttention requires up_v_weight for latent-to-value projection."
+            )
+        if value is not None:
+            raise RuntimeError(
+                "Absorbed DSAttention expects value=None (latent path). "
+                "Received absorbed layout with explicit value tensor."
+            )
+        output = None
+        if dsa_kernels.use_fused_dsa_kernels(config):
+            output = dsa_kernels.run_fused_absorbed_sparse_attention(
+                config, query, key, topk_indices, softmax_scale, latent_v_channels
+            )
+        if output is None:
+            output = _unfused_absorbed_dsa_fn(
+                query,
+                key,
+                topk_indices,
+                softmax_scale,
+                latent_v_channels,
+                mask=mask,
+                varlen_starts=varlen_starts,
+                varlen_ends=varlen_ends,
+                key_positions=key_positions,
+            )
+        output = torch.einsum("sbhc,hdc->sbhd", output, up_v_weight).contiguous()
+        output = output.view(output.size(0), output.size(1), -1)
+        return output
+
+    return unfused_dsa_fn(
+        query,
+        key,
+        value,
+        topk_indices,
+        softmax_scale,
+        mask=mask,
+        varlen_starts=varlen_starts,
+        varlen_ends=varlen_ends,
+        key_positions=key_positions,
+    )
+
+
+def _normalize_dsattention_output_rank(output: torch.Tensor, target_ndim: int) -> torch.Tensor:
+    """Normalize DSAttention output rank to match caller hidden-state rank."""
+    if target_ndim not in (2, 3):
+        raise RuntimeError(f"DSAttention expected x.ndim in (2, 3), got {target_ndim}")
+
+    if output.ndim == 4:
+        output = output.reshape(output.size(0), output.size(1), -1)
+    elif output.ndim not in (2, 3):
+        raise RuntimeError(
+            f"DSAttention produced unexpected output rank {output.ndim}; expected 2D/3D/4D."
+        )
+
+    if target_ndim == 3 and output.ndim == 2:
+        output = output.unsqueeze(1)
+    elif target_ndim == 2 and output.ndim == 3:
+        if output.size(1) != 1:
+            raise RuntimeError(
+                "DSAttention cannot squeeze non-singleton batch dim for packed output: "
+                f"shape={tuple(output.shape)}"
+            )
+        output = output.squeeze(1)
+
+    if output.ndim != target_ndim:
+        raise RuntimeError(
+            "DSAttention output rank mismatch after normalization: "
+            f"target_ndim={target_ndim}, output_shape={tuple(output.shape)}"
+        )
+    return output
+
+
+def _validate_nonpacked_cp_uniform_length(
+    sq: int,
+    skv: int,
+    cp_size: int,
+    cp_group: Optional[torch.distributed.ProcessGroup],
+    device: torch.device,
+) -> None:
+    """Validate the uniform-length precondition for non-packed allgather CP."""
+    expected_skv = sq * cp_size
+    if (
+        cp_group is not None
+        and torch.distributed.is_available()
+        and torch.distributed.is_initialized()
+        and cp_group.size() == cp_size
+    ):
+        local_len = torch.tensor([sq], device=device, dtype=torch.int64)
+        all_lens = [torch.empty_like(local_len) for _ in range(cp_size)]
+        torch.distributed.all_gather(all_lens, local_len, group=cp_group)
+        all_lens = torch.cat(all_lens)
+        if not torch.all(all_lens == sq):
+            raise RuntimeError(
+                "Non-packed DSA allgather CP expects uniform per-rank sequence lengths; "
+                f"got per-rank lengths {all_lens.tolist()}."
+            )
+        expected_skv = int(all_lens.sum().item())
+
+    if skv != sq and skv != expected_skv:
+        raise RuntimeError(
+            "Non-packed DSA allgather CP expects uniform per-rank sequence lengths; "
+            f"got local query length {sq} and key length {skv} for cp_size={cp_size}."
+        )
+
+
 def rotate_activation(x: torch.Tensor) -> torch.Tensor:
     """Apply Hadamard rotation activation.
     Reference:
@@ -167,6 +353,12 @@ def compute_dsa_indexer_loss(
     loss_coeff: float,
     sparse_loss: bool,
     pg_collection: ProcessGroupCollection,
+    mask: Optional[torch.Tensor] = None,
+    varlen_starts: Optional[torch.Tensor] = None,
+    varlen_ends: Optional[torch.Tensor] = None,
+    key_positions: Optional[torch.Tensor] = None,
+    query_valid_rows: Optional[torch.Tensor] = None,
+    calculate_per_token_loss: bool = False,
 ) -> torch.Tensor:
     """
     Compute KL divergence loss between index_scores and true attention_scores.
@@ -187,12 +379,23 @@ def compute_dsa_indexer_loss(
         sparse_loss: bool, whether to use sparse indexer loss. If True, only the topk
             indices will be used to compute the loss.
         pg_collection: Process group collection, must have TP process group.
+        mask: Optional additive attention mask. Supports shape [sq, sk] or [b, sq, sk].
+            Invalid positions should be -inf.
+        varlen_starts: Optional row-wise key start bounds [sq] for packed THD.
+        varlen_ends: Optional row-wise key end bounds [sq] for packed THD.
+        key_positions: Optional global key positions [sk] for packed THD.
 
     Returns:
         index_loss: KL divergence loss (scalar).
     """
+    query, _ = dsa_layout.ensure_sbhd(query, "query")
+    key, _ = dsa_layout.ensure_sbhd(key, "key")
+
     sq, b, np, hn = query.size()
     sk = key.size(0)
+    query_valid_rows = dsa_masking.normalize_query_valid_rows(
+        query_valid_rows, b=b, sq=sq, device=index_scores.device
+    )
 
     # [sq, b, np, hn] -> [b, np, sq, hn] -> [b * np, sq, hn]
     query = query.permute(1, 2, 0, 3).reshape(b * np, sq, hn)
@@ -202,29 +405,58 @@ def compute_dsa_indexer_loss(
     attention_scores = torch.bmm(query.float(), key.float()) * softmax_scale
     # Reshape to [b, np, sq, sk]
     attention_scores = attention_scores.reshape(b, np, sq, sk)
-
-    # causal_mask [sq, sk]
-    causal_mask = torch.triu(
-        torch.full((sq, sk), float('-inf'), dtype=torch.float32, device=attention_scores.device),
-        diagonal=1,
+    varlen_starts, varlen_ends, key_positions = dsa_masking.normalize_varlen_bounds(
+        mask=mask,
+        varlen_starts=varlen_starts,
+        varlen_ends=varlen_ends,
+        key_positions=key_positions,
+        sk=sk,
+        device=attention_scores.device,
     )
+
+    if varlen_starts is not None:
+        attention_scores = dsa_masking.apply_starts_ends_mask_to_scores(
+            attention_scores, varlen_starts, varlen_ends, key_positions
+        )
+        index_scores = dsa_masking.apply_starts_ends_mask_to_scores(
+            index_scores, varlen_starts, varlen_ends, key_positions
+        )
+        base_valid_mask = (
+            dsa_masking.build_valid_mask_from_starts_ends(varlen_starts, varlen_ends, key_positions)
+            .unsqueeze(0)
+            .expand(b, sq, sk)
+        )
+    else:
+        _, attn_score_mask, index_score_mask, base_valid_mask = dsa_masking.prepare_additive_mask(
+            mask, sq=sq, sk=sk, b=b, device=attention_scores.device
+        )
+        # [b, np, sq, sk] + [1/b, 1, sq, sk] -> [b, np, sq, sk]
+        attention_scores += attn_score_mask
+        # [b, sq, sk] + [1/b, sq, sk] -> [b, sq, sk]
+        index_scores += index_score_mask
+
     # index_mask [b, sq, sk]
     index_mask = torch.full(
-        (b, sq, sk), float("-inf"), dtype=torch.float32, device=causal_mask.device
-    ).scatter_(-1, topk_indices, 0)
+        (b, sq, sk), float("-inf"), dtype=torch.float32, device=attention_scores.device
+    )
+    dsa_masking.scatter_topk_into_index_mask(index_mask, topk_indices, seq_chunk_size=256)
 
-    # [b, np, sq, skv] + [1, 1, sq, skv] -> [b, np, sq, skv]
-    attention_scores += causal_mask.view(1, 1, sq, sk)
     if sparse_loss:
         # [b, np, sq, sk] + [b, 1, sq, sk] -> [b, np, sq, sk]
         attention_scores += index_mask.view(b, 1, sq, sk)
         # [b, sq, sk] + [b, sq, sk] -> [b, sq, sk]
         index_scores += index_mask
+        index_valid_mask = base_valid_mask & (index_mask == 0)
+    else:
+        index_valid_mask = base_valid_mask
+    attention_valid_mask = index_valid_mask if sparse_loss else base_valid_mask
 
     # [b, np, sq, sk] -> [b, np, sq, sk]
-    attention_scores = torch.nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32)
+    attention_scores = dsa_masking.masked_softmax(
+        attention_scores.float(), attention_valid_mask.unsqueeze(1).expand(b, np, sq, sk), dim=-1
+    )
     # [b, sq, sk] -> [b, sq, sk]
-    index_scores = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32)
+    index_scores = dsa_masking.masked_softmax(index_scores.float(), index_valid_mask, dim=-1)
 
     # Sum attention scores across heads.
     # [batch, heads, seqlen_q, seqlen_k] -> [batch, seqlen_q, seqlen_k]
@@ -234,7 +466,9 @@ def compute_dsa_indexer_loss(
         torch.distributed.all_reduce(attention_scores.contiguous(), group=pg_collection.tp)
     # L1 normalize target on the last dimension. Doesn't use abs() because attention_scores are
     # obtained from softmax so they are already non-negative.
-    attention_scores = attention_scores / attention_scores.sum(dim=-1, keepdim=True)
+    attention_scores = attention_scores / attention_scores.sum(dim=-1, keepdim=True).clamp_min(
+        1e-10
+    )
 
     # Compute KL divergence: KL(target || index) = target(x) * log(target(x) / index(x))
     # kl_per_element [b, sq, sk]
@@ -243,8 +477,19 @@ def compute_dsa_indexer_loss(
     )
 
     # [b, sq, sk] -> [b, sq] -> [1]
-    # Each token has same weight in the loss.
-    kl_div = kl_per_element.sum(dim=-1).mean()
+    # Each real token has the same weight in the loss.
+    kl_per_row = kl_per_element.sum(dim=-1)
+    if calculate_per_token_loss:
+        if query_valid_rows is None:
+            kl_div = kl_per_row.sum()
+        else:
+            kl_div = (kl_per_row * query_valid_rows.to(dtype=torch.float32)).sum()
+    elif query_valid_rows is None:
+        kl_div = kl_per_row.mean()
+    else:
+        valid_row_count = query_valid_rows.sum().to(dtype=torch.float32, device=kl_per_row.device)
+        valid_row_count = valid_row_count.clamp_min(1.0)
+        kl_div = (kl_per_row * query_valid_rows.to(dtype=torch.float32)).sum() / valid_row_count
 
     # Scale by coefficient.
     indexer_loss = kl_div * loss_coeff
@@ -252,7 +497,9 @@ def compute_dsa_indexer_loss(
     return indexer_loss
 
 
-def _compute_index_scores(q: torch.Tensor, weights: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
+def _compute_index_scores(
+    q: torch.Tensor, weights: torch.Tensor, k: torch.Tensor, use_relu: bool = True
+) -> torch.Tensor:
     """
     Perform index score using BF16 precision.
 
@@ -260,7 +507,7 @@ def _compute_index_scores(q: torch.Tensor, weights: torch.Tensor, k: torch.Tenso
         https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/kernel.py#L254-L274
     This is a BF16 implementation of the `fp8_index` logic:
         1. Compute attention scores: q @ k^T;
-        2. Apply ReLU activation;
+        2. Optionally apply ReLU activation (DeepSeek V3.2 only; disabled for GLM5);
         3. Weight by attention weights;
         4. Sum across attention heads.
 
@@ -277,8 +524,9 @@ def _compute_index_scores(q: torch.Tensor, weights: torch.Tensor, k: torch.Tenso
     #   -> [seqlen_q, batch, index_n_heads, seqlen_k]
     index_scores = torch.einsum('sbhd,tbd->sbht', q.float(), k.float())
 
-    # Apply ReLU activation.
-    index_scores = torch.relu(index_scores)
+    # 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]
@@ -301,33 +549,80 @@ def fused_qk_topk_naive(
     weights: torch.Tensor,
     index_topk: int,
     mask: Optional[torch.Tensor] = None,
+    varlen_starts: Optional[torch.Tensor] = None,
+    varlen_ends: Optional[torch.Tensor] = None,
+    key_positions: Optional[torch.Tensor] = None,
+    use_relu: bool = True,
 ):
     """Naive implementation of QK Topk."""
-    seqlen = q.size(0)
+    sk = k.size(0)
     # =========================================
     # Compute index scores
     # =========================================
     # [batch, seqlen, seqlen]
-    index_scores = _compute_index_scores(q, weights, k)
-    if mask is not None:
+    index_scores = _compute_index_scores(q, weights, k, use_relu=use_relu)
+    varlen_starts, varlen_ends, key_positions = dsa_masking.normalize_varlen_bounds(
+        mask=mask,
+        varlen_starts=varlen_starts,
+        varlen_ends=varlen_ends,
+        key_positions=key_positions,
+        sk=sk,
+        device=index_scores.device,
+    )
+    if varlen_starts is not None:
+        index_scores = dsa_masking.apply_starts_ends_mask_to_scores(
+            index_scores, varlen_starts, varlen_ends, key_positions
+        )
+    elif mask is not None:
         assert mask.dtype == index_scores.dtype, "Mask dtype must match index scores dtype"
         index_scores = index_scores + mask
 
     # =========================================
     # Select top-k indices
     # =========================================
-    topk_k = min(index_topk, seqlen)
-    # [batch, seqlen, index_topk]
-    topk_indices = index_scores.topk(topk_k, dim=-1)[1]
+    topk_k = min(index_topk, sk)
+    if topk_k > 0:
+        topk_scores, topk_indices = index_scores.topk(topk_k, dim=-1)
+        topk_indices = topk_indices.masked_fill(topk_scores == float("-inf"), -1)
+    else:
+        topk_indices = torch.empty(
+            index_scores.shape[:-1] + (0,), dtype=torch.int64, device=index_scores.device
+        )
 
     return index_scores, topk_indices
 
 
 def fwd_fused_indexer_loss_naive(
-    q, weights, k, query, key, topk, softmax_scale, loss_coeff, mask, sparse_loss, pg_collection
+    q,
+    weights,
+    k,
+    query,
+    key,
+    topk,
+    softmax_scale,
+    loss_coeff,
+    mask,
+    sparse_loss,
+    pg_collection,
+    varlen_starts=None,
+    varlen_ends=None,
+    key_positions=None,
+    query_valid_rows=None,
+    calculate_per_token_loss: bool = False,
+    use_relu: bool = True,
 ):
     """Naive implementation of forward pass for indexer loss."""
-    index_scores, topk_indices = fused_qk_topk_naive(q, k, weights, topk, mask)
+    index_scores, topk_indices = fused_qk_topk_naive(
+        q,
+        k,
+        weights,
+        topk,
+        mask=mask,
+        varlen_starts=varlen_starts,
+        varlen_ends=varlen_ends,
+        key_positions=key_positions,
+        use_relu=use_relu,
+    )
 
     indexer_loss = compute_dsa_indexer_loss(
         index_scores,
@@ -338,6 +633,12 @@ def fwd_fused_indexer_loss_naive(
         loss_coeff,
         sparse_loss,
         pg_collection,
+        mask=mask,
+        varlen_starts=varlen_starts,
+        varlen_ends=varlen_ends,
+        key_positions=key_positions,
+        query_valid_rows=query_valid_rows,
+        calculate_per_token_loss=calculate_per_token_loss,
     )
 
     return topk_indices, indexer_loss
@@ -353,14 +654,27 @@ def bwd_fused_indexer_loss_naive(
     softmax_scale,
     loss_coeff,
     sparse_loss,
+    mask,
     grad_loss,
     pg_collection,
+    varlen_starts=None,
+    varlen_ends=None,
+    key_positions=None,
+    query_valid_rows=None,
+    calculate_per_token_loss: bool = False,
+    use_relu: bool = True,
 ):
     """Naive implementation of backward pass for indexer loss."""
-    index_scores = _compute_index_scores(q, weights, k)  # [B, Sq, Sk]
+    query, _ = dsa_layout.ensure_sbhd(query, "query")
+    key, _ = dsa_layout.ensure_sbhd(key, "key")
+
+    index_scores = _compute_index_scores(q, weights, k, use_relu=use_relu)  # [B, Sq, Sk]
 
     sq, b, np, hn = query.size()
     sk = key.size(0)
+    query_valid_rows = dsa_masking.normalize_query_valid_rows(
+        query_valid_rows, b=b, sq=sq, device=query.device
+    )
 
     # [sq, b, np, hn] -> [b, np, sq, hn] -> [b * np, sq, hn]
     query_reshaped = query.permute(1, 2, 0, 3).reshape(b * np, sq, hn)
@@ -373,24 +687,41 @@ def bwd_fused_indexer_loss_naive(
 
     # Reshape to [b, np, sq, sk]
     attention_scores = attention_scores.reshape(b, np, sq, sk)
-
-    # causal_mask [sq, sk]
-    causal_mask = torch.triu(
-        torch.full((sq, sk), float('-inf'), dtype=torch.float32, device=attention_scores.device),
-        diagonal=1,
+    varlen_starts, varlen_ends, key_positions = dsa_masking.normalize_varlen_bounds(
+        mask=mask,
+        varlen_starts=varlen_starts,
+        varlen_ends=varlen_ends,
+        key_positions=key_positions,
+        sk=sk,
+        device=attention_scores.device,
     )
+
+    if varlen_starts is not None:
+        attention_scores = dsa_masking.apply_starts_ends_mask_to_scores(
+            attention_scores, varlen_starts, varlen_ends, key_positions
+        )
+        index_scores = dsa_masking.apply_starts_ends_mask_to_scores(
+            index_scores, varlen_starts, varlen_ends, key_positions
+        )
+        base_valid_mask = (
+            dsa_masking.build_valid_mask_from_starts_ends(varlen_starts, varlen_ends, key_positions)
+            .unsqueeze(0)
+            .expand(b, sq, sk)
+        )
+    else:
+        _, attn_score_mask, index_score_mask, base_valid_mask = dsa_masking.prepare_additive_mask(
+            mask, sq=sq, sk=sk, b=b, device=attention_scores.device
+        )
+        # [b, np, sq, sk] + [1/b, 1, sq, sk] -> [b, np, sq, sk]
+        attention_scores = attention_scores + attn_score_mask
+        # [b, sq, sk] + [1/b, sq, sk] -> [b, sq, sk]
+        index_scores = index_scores + index_score_mask
+
     # index_mask [b, sq, sk]
     index_mask = torch.full(
-        (b, sq, sk), float("-inf"), dtype=torch.float32, device=causal_mask.device
-    ).scatter_(-1, topk_indices, 0)
-
-    # Apply causal mask to both attention and index scores
-    # [b, np, sq, skv] + [1, 1, sq, skv] -> [b, np, sq, skv]
-    attention_scores = attention_scores + causal_mask.view(1, 1, sq, sk)
-    # [b, sq, sk] + [1, sq, sk] -> [b, sq, sk]
-    index_scores = index_scores + causal_mask.unsqueeze(0)
-    # Free causal_mask - no longer needed
-    del causal_mask
+        (b, sq, sk), float("-inf"), dtype=torch.float32, device=attention_scores.device
+    )
+    dsa_masking.scatter_topk_into_index_mask(index_mask, topk_indices, seq_chunk_size=256)
 
     if sparse_loss:
         # [b, np, sq, sk] + [b, 1, sq, sk] -> [b, np, sq, sk]
@@ -398,14 +729,21 @@ def bwd_fused_indexer_loss_naive(
         # [b, sq, sk] + [b, sq, sk] -> [b, sq, sk]
         index_scores = index_scores + index_mask
 
-    # Compute softmax for both
-    attention_scores_softmax = torch.nn.functional.softmax(
-        attention_scores, dim=-1, dtype=torch.float32
+    # Compute softmax for both.
+    if sparse_loss:
+        index_valid_mask = base_valid_mask & (index_mask == 0)
+    else:
+        index_valid_mask = base_valid_mask
+    attention_valid_mask = index_valid_mask if sparse_loss else base_valid_mask
+    attention_scores_softmax = dsa_masking.masked_softmax(
+        attention_scores.float(), attention_valid_mask.unsqueeze(1).expand(b, np, sq, sk), dim=-1
     )
     # Free attention_scores immediately
     del attention_scores
 
-    index_scores_softmax = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32)
+    index_scores_softmax = dsa_masking.masked_softmax(
+        index_scores.float(), index_valid_mask, dim=-1
+    )
     # Free index_scores - no longer needed after softmax
     del index_scores
 
@@ -421,7 +759,7 @@ def bwd_fused_indexer_loss_naive(
     # L1 normalize
     attention_scores_normalized = attention_scores_sum / attention_scores_sum.sum(
         dim=-1, keepdim=True
-    )
+    ).clamp_min(1e-10)
     # Free attention_scores_sum - no longer needed after normalization
     del attention_scores_sum
 
@@ -429,12 +767,27 @@ def bwd_fused_indexer_loss_naive(
     # where kl_div = kl_per_element.sum(dim=-1).mean()
     grad_kl_div = grad_loss * loss_coeff  # scalar
 
-    # Backward through mean: distribute gradient equally
-    grad_kl_per_row = grad_kl_div / (b * sq)  # scalar value for each row
+    if calculate_per_token_loss:
+        grad_kl_per_row = grad_kl_div
+    else:
+        valid_row_count = (
+            query_valid_rows.sum().to(
+                dtype=torch.float32, device=attention_scores_normalized.device
+            )
+            if query_valid_rows is not None
+            else torch.tensor(
+                float(b * sq), dtype=torch.float32, device=attention_scores_normalized.device
+            )
+        ).clamp_min(1.0)
+        grad_kl_per_row = grad_kl_div / valid_row_count  # scalar value for each real row
 
     # Backward through sum(dim=-1): broadcast back to [b, sq, sk]
     # Each element in a row contributes to the sum, so gradient is same for all
     grad_kl_per_element = grad_kl_per_row.view(1, 1, 1).expand(b, sq, sk)
+    if query_valid_rows is not None:
+        grad_kl_per_element = grad_kl_per_element * query_valid_rows.unsqueeze(-1).to(
+            dtype=grad_kl_per_element.dtype
+        )
 
     # Backward through kl_per_element = target * (log(target) - log(index))
     # ∂kl/∂index_softmax = -target / index_softmax
@@ -450,22 +803,18 @@ def bwd_fused_indexer_loss_naive(
     # Free intermediate tensors
     del index_scores_softmax, grad_index_scores_softmax, sum_grad
 
-    # Zero out gradients for masked positions
-    # Create a mask for valid (non-masked) positions
-    # Causal mask: position (i, j) is valid if j <= i
-    causal_valid_mask = torch.tril(
-        torch.ones((sq, sk), device=q.device, dtype=torch.bool)
-    )  # [sq, sk]
+    # Zero out gradients for masked positions.
     if sparse_loss:
-        # Also apply index mask - only topk positions are valid
-        index_valid_mask = index_mask == 0  # [b, sq, sk]
-        del index_mask  # Free index_mask immediately after use
-        valid_mask = causal_valid_mask.unsqueeze(0) & index_valid_mask  # [b, sq, sk]
+        # Also apply index mask - only topk positions are valid.
+        del index_mask
+        valid_mask = base_valid_mask & index_valid_mask  # [b, sq, sk]
         del index_valid_mask
     else:
-        del index_mask  # Free index_mask even if not used for sparse_loss
-        valid_mask = causal_valid_mask.unsqueeze(0).expand(b, sq, sk)  # [b, sq, sk]
-    del causal_valid_mask
+        del index_mask
+        valid_mask = base_valid_mask  # [b, sq, sk]
+    del base_valid_mask
+    if query_valid_rows is not None:
+        valid_mask = valid_mask & query_valid_rows.unsqueeze(-1)
 
     grad_index_scores_logits = grad_index_scores_logits * valid_mask.float()
     del valid_mask
@@ -480,22 +829,27 @@ def bwd_fused_indexer_loss_naive(
 
     # Compute forward values needed for backward
     scores = torch.einsum('sbhd,tbd->sbht', q.float(), k.float())  # [sq, b, h, sk]
-    # Compute relu_mask before relu (saves memory vs keeping both scores and relu output)
-    relu_mask = scores > 0
-    scores_after_relu = torch.relu(scores)
+
+    # 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
 
-    # Backward through multiplication by weights: index_scores_per_head * weights
-    # ∂L/∂weights = grad * relu_scores (sum over sk)
-    grad_weights = (grad_weighted_scores * scores_after_relu).sum(dim=-1)  # [sq, b, h]
+    # ∂L/∂weights = grad * scores_for_weights (sum over sk)
+    grad_weights = (grad_weighted_scores * scores_for_weights).sum(dim=-1)  # [sq, b, h]
 
-    # ∂L/∂relu_scores = grad * weights
-    grad_scores_after_relu = grad_weighted_scores * weights.unsqueeze(-1)  # [sq, b, h, sk]
-    del grad_weighted_scores, scores_after_relu
+    # ∂L/∂scores = grad * weights
+    grad_scores = grad_weighted_scores * weights.unsqueeze(-1)  # [sq, b, h, sk]
+    del grad_weighted_scores, scores_for_weights
 
-    # Backward through ReLU
-    grad_scores = grad_scores_after_relu * relu_mask.float()  # [sq, b, h, sk]
-    del grad_scores_after_relu, relu_mask
+    # Backward through ReLU (skip when use_relu=False)
+    if use_relu:
+        grad_scores = grad_scores * relu_mask.float()
+        del relu_mask
 
     # Backward through einsum 'sbhd,tbd->sbht'
     # ∂L/∂q = einsum('sbht,tbd->sbhd', grad_scores, k)
@@ -507,6 +861,27 @@ def bwd_fused_indexer_loss_naive(
     return grad_q.to(q.dtype), grad_weights.to(weights.dtype), grad_k.to(k.dtype)
 
 
+_FUSED_DSA_INDEXER_LOSS_INPUT_NAMES = (
+    "q",
+    "weights",
+    "k",
+    "query",
+    "key",
+    "softmax_scale",
+    "topk",
+    "loss_coeff",
+    "mask",
+    "sparse_loss",
+    "pg_collection",
+    "varlen_starts",
+    "varlen_ends",
+    "key_positions",
+    "query_valid_rows",
+    "calculate_per_token_loss",
+    "use_relu",
+)
+
+
 class FusedDSAIndexerLoss(torch.autograd.Function):
     """Fused implementation of DSA Indexer Loss."""
 
@@ -524,6 +899,12 @@ def forward(
         mask,
         sparse_loss,
         pg_collection,
+        varlen_starts=None,
+        varlen_ends=None,
+        key_positions=None,
+        query_valid_rows=None,
+        calculate_per_token_loss: bool = False,
+        use_relu: bool = True,
     ):
         """
         Fused forward: index_scores never materialized in full.
@@ -540,6 +921,12 @@ def forward(
             mask,
             sparse_loss,
             pg_collection,
+            varlen_starts=varlen_starts,
+            varlen_ends=varlen_ends,
+            key_positions=key_positions,
+            query_valid_rows=query_valid_rows,
+            calculate_per_token_loss=calculate_per_token_loss,
+            use_relu=use_relu,
         )
 
         # Save for backward (recomputation strategy)
@@ -547,7 +934,14 @@ def forward(
         ctx.softmax_scale = softmax_scale
         ctx.loss_coeff = loss_coeff
         ctx.sparse_loss = sparse_loss
+        ctx.mask = mask
         ctx.pg_collection = pg_collection
+        ctx.varlen_starts = varlen_starts
+        ctx.varlen_ends = varlen_ends
+        ctx.key_positions = key_positions
+        ctx.query_valid_rows = query_valid_rows
+        ctx.calculate_per_token_loss = calculate_per_token_loss
+        ctx.use_relu = use_relu
 
         return topk_indices, loss
 
@@ -568,12 +962,26 @@ def backward(ctx, grad_topk_indices, grad_loss):
             ctx.softmax_scale,
             ctx.loss_coeff,
             ctx.sparse_loss,
+            ctx.mask,
             grad_loss,
             ctx.pg_collection,
+            varlen_starts=ctx.varlen_starts,
+            varlen_ends=ctx.varlen_ends,
+            key_positions=ctx.key_positions,
+            query_valid_rows=ctx.query_valid_rows,
+            calculate_per_token_loss=ctx.calculate_per_token_loss,
+            use_relu=ctx.use_relu,
         )
 
-        # query and key are detached in forward, so return None for their gradients
-        return grad_q, grad_weights, grad_k, None, None, None, None, None, None, None, None
+        grad_by_name = {
+            "q": grad_q,
+            "weights": grad_weights,
+            "k": grad_k,
+            # query and key are detached in forward, so return None for their gradients.
+            "query": None,
+            "key": None,
+        }
+        return tuple(grad_by_name.get(name) for name in _FUSED_DSA_INDEXER_LOSS_INPUT_NAMES)
 
 
 class DSAIndexerLossAutoScaler(torch.autograd.Function):
@@ -583,7 +991,7 @@ class DSAIndexerLossAutoScaler(torch.autograd.Function):
     to train the indexer to predict attention scores without affecting the forward pass.
     """
 
-    main_loss_backward_scale: torch.Tensor = None
+    main_loss_backward_scale: Optional[torch.Tensor] = None
 
     @staticmethod
     def forward(ctx, output: torch.Tensor, indexer_loss: torch.Tensor):
@@ -615,7 +1023,9 @@ def backward(ctx, grad_output: torch.Tensor):
             DSAIndexerLossAutoScaler.main_loss_backward_scale = torch.tensor(
                 1.0, device=indexer_loss.device
             )
-        indexer_loss_backward_scale = DSAIndexerLossAutoScaler.main_loss_backward_scale
+        indexer_loss_backward_scale = DSAIndexerLossAutoScaler.main_loss_backward_scale.to(
+            device=indexer_loss.device
+        )
         scaled_indexer_loss_grad = torch.ones_like(indexer_loss) * indexer_loss_backward_scale
         return grad_output, scaled_indexer_loss_grad
 
@@ -626,6 +1036,10 @@ def set_loss_scale(scale: torch.Tensor):
         Args:
             scale: The scale value to set.
         """
+        if not isinstance(scale, torch.Tensor):
+            raise TypeError("DSAIndexerLossAutoScaler.set_loss_scale requires a torch.Tensor.")
+        scale = scale.detach()
+
         if DSAIndexerLossAutoScaler.main_loss_backward_scale is None:
             DSAIndexerLossAutoScaler.main_loss_backward_scale = scale
         else:
@@ -757,11 +1171,13 @@ def __init__(
 
         k_norm_config = copy.copy(self.config)
         k_norm_config.normalization = "LayerNorm"
+        k_norm_eps = (
+            self.config.dsa_indexer_k_norm_epsilon
+            if self.config.dsa_indexer_k_norm_epsilon is not None
+            else self.config.layernorm_epsilon
+        )
         self.k_norm = build_module(
-            submodules.k_norm,
-            config=k_norm_config,
-            hidden_size=self.index_head_dim,
-            eps=self.config.layernorm_epsilon,
+            submodules.k_norm, config=k_norm_config, hidden_size=self.index_head_dim, eps=k_norm_eps
         )
 
         self.linear_weights_proj = build_module(
@@ -776,7 +1192,13 @@ def __init__(
             parallel_mode="duplicated",
         )
 
-    def _apply_rope(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor, mscale: float):
+    def _apply_rope(
+        self,
+        x: torch.Tensor,
+        rotary_pos_emb: torch.Tensor,
+        mscale: float,
+        cu_seqlens: Optional[torch.Tensor] = None,
+    ):
         """Apply RoPE to the input tensor."""
         # x_pe   [seqlen, batch, *, qk_pos_emb_head_dim]
         # x_nope [seqlen, batch, *, index_head_dim - qk_pos_emb_head_dim]
@@ -785,17 +1207,25 @@ def _apply_rope(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor, mscale: flo
         x_pe, x_nope = torch.split(
             x, [self.qk_pos_emb_head_dim, self.index_head_dim - self.qk_pos_emb_head_dim], dim=-1
         )
+        squeezed_batch_dim = False
+        if cu_seqlens is not None and cu_seqlens.device != x_pe.device:
+            cu_seqlens = cu_seqlens.to(device=x_pe.device)
+        # THD RoPE path expects [t, h, d], while indexer tensors are [t, 1, h, d].
+        if cu_seqlens is not None and x_pe.ndim == 4 and x_pe.size(1) == 1:
+            x_pe = x_pe.squeeze(1)
+            squeezed_batch_dim = True
         x_pe = apply_rotary_pos_emb(
             x_pe,
             rotary_pos_emb,
             config=self.config,
-            cu_seqlens=None,
+            cu_seqlens=cu_seqlens,
             mscale=mscale,
             cp_group=self.pg_collection.cp,
             # This flag is for the MLA-style interleaving in RoPE.
-            # Set it to False, as indexer does not apply interleaved RoPE.
-            mla_rotary_interleaved=False,
+            mla_rotary_interleaved=self.config.dsa_indexer_rope_interleaved,
         )
+        if squeezed_batch_dim:
+            x_pe = x_pe.unsqueeze(1)
         # [seqlen, batch, *, index_head_dim]
         x = torch.cat([x_pe, x_nope], dim=-1)
         return x
@@ -804,6 +1234,8 @@ def forward_before_topk(
         self, x: torch.Tensor, qr: torch.Tensor, packed_seq_params: Optional[PackedSeqParams] = None
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         """All computations before topk."""
+        packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == "thd"
+
         # =========================================
         # Prepare RoPE params
         # =========================================
@@ -811,10 +1243,14 @@ def forward_before_topk(
             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=False)
+            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=False)
+            rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq)
+        if packed_seq:
+            cu_seqlens_q, cu_seqlens_kv = dsa_layout.get_packed_qk_cu_seqlens(packed_seq_params)
+        else:
+            cu_seqlens_q = cu_seqlens_kv = None
 
         # =========================================
         # Gather inputs if sp is enabled
@@ -836,25 +1272,30 @@ def forward_before_topk(
         # [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)
-        q = self._apply_rope(q, rotary_pos_emb, mscale)
+        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)
-        k = self.k_norm(k)
+        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)
+        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
         # =========================================
-        q = rotate_activation(q)
-        k = rotate_activation(k)
+        if self.config.dsa_indexer_rotate_activation:
+            q = rotate_activation(q)
+            k = rotate_activation(k)
 
         # =========================================
         # Prepare weights for index scores
@@ -880,22 +1321,23 @@ def forward_with_scores(
         Args:
             x: hidden states [seqlen, batch, hidden_size].
             qr: Low-rank query tensor [seqlen, batch, q_lora_rank].
-            mask: Attention mask [batch, seqlen, seqlen].
+            mask: Optional additive attention mask [seqlen, seqlen] or
+                [batch, seqlen, seqlen].
             packed_seq_params: Packed sequence parameters for variable length sequences.
 
         Returns:
             index_scores: Index scores [batch, seqlen, seqlen].
             topk_indices: Top-k indices [batch, seqlen, index_topk].
         """
-        assert packed_seq_params is None, "Packed sequence is not supported for DSAttention"
-
         # [seqlen, batch, index_n_heads * index_head_dim]
         # [seqlen, batch, index_head_dim]
         # [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)
+        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
 
@@ -922,56 +1364,151 @@ def forward(
         return topk_indices
 
 
-def unfused_dsa_fn(query, key, value, topk_indices, softmax_scale):
+def unfused_dsa_fn(
+    query,
+    key,
+    value,
+    topk_indices,
+    softmax_scale,
+    mask: Optional[torch.Tensor] = None,
+    varlen_starts: Optional[torch.Tensor] = None,
+    varlen_ends: Optional[torch.Tensor] = None,
+    key_positions: Optional[torch.Tensor] = None,
+):
     """
     Unfused sparse attention implementation.
+
+    This path uses chunked sparse softmax accumulation over top-k selected keys
+    to avoid materializing full [b, np, sq, skv] attention score tensors.
     """
+    if value is None:
+        raise NotImplementedError("DSAttention unfused path requires value tensor.")
+
+    query, query_was_thd = dsa_layout.ensure_sbhd(query, "query")
+    key, _ = dsa_layout.ensure_sbhd(key, "key")
+    value, _ = dsa_layout.ensure_sbhd(value, "value")
+
     sq, b, np, hn = query.size()
     skv = key.size(0)
+    nk = key.size(2)
     hnv = value.size(3)
+    nv = value.size(2)
+
+    # [sq, b, np, hn] -> [b, np, sq, hn]
+    query_b = query.permute(1, 2, 0, 3).contiguous()
+    # [skv, b, nk, hn] -> [b, nk, skv, hn]
+    key_b = key.permute(1, 2, 0, 3).contiguous()
+    # [skv, b, nv, hnv] -> [b, nv, skv, hnv]
+    value_b = value.permute(1, 2, 0, 3).contiguous()
+    if nk == 1 and np > 1:
+        key_b = key_b.expand(b, np, skv, hn)
+    else:
+        assert nk == np, "key head count must be 1 (MQA) or match query heads"
+    if nv == 1 and np > 1:
+        value_b = value_b.expand(b, np, skv, hnv)
+    else:
+        assert nv == np, "value head count must be 1 (MQA) or match query heads"
+
+    row_mask, varlen_starts, varlen_ends, key_positions = dsa_masking.prepare_sparse_mask_context(
+        mask=mask,
+        varlen_starts=varlen_starts,
+        varlen_ends=varlen_ends,
+        key_positions=key_positions,
+        sq=sq,
+        sk=skv,
+        b=b,
+        device=query.device,
+    )
 
-    # ===================================
-    # Raw attention scores [b, np, sq, skv]
-    # ===================================
-    # [sq, b, np, hn] -> [b, np, sq, hn] -> [b * np, sq, hn]
-    query = query.permute(1, 2, 0, 3).reshape(b * np, sq, hn)
-    # [skv, b, np, hn] -> [b, np, hn, skv] -> [b * np, hn, skv]
-    key = key.permute(1, 2, 3, 0).reshape(b * np, hn, skv)
-    # Compute attention scores [b * np, sq, skv]
-    attention_scores = torch.bmm(query.float(), key.float()) * softmax_scale
-    # Reshape to [b, np, sq, skv]
-    attention_scores = attention_scores.reshape(b, np, sq, skv)
+    seq_chunk_size = 512
+    head_chunk_size = 16
+    topk_chunk_size = 1024
+    safe_k_max = max(0, skv - 1)
+    output = torch.empty((sq, b, np * hnv), dtype=value.dtype, device=query.device)
+
+    for bi in range(b):
+        for h0 in range(0, np, head_chunk_size):
+            h1 = min(h0 + head_chunk_size, np)
+            h_chunk = h1 - h0
+            out_h0 = h0 * hnv
+            out_h1 = h1 * hnv
+            k_chunk = key_b[bi, h0:h1, :, :].contiguous()  # [h_chunk, skv, hn]
+            v_chunk = value_b[bi, h0:h1, :, :].contiguous()  # [h_chunk, skv, hnv]
+            flat_k = k_chunk.reshape(h_chunk * skv, hn)
+            flat_v = v_chunk.reshape(h_chunk * skv, hnv)
+            head_offsets = (
+                torch.arange(h_chunk, device=query.device, dtype=torch.int64).view(-1, 1, 1) * skv
+            )
 
-    # ===================================
-    # Apply sparse mask from indexer
-    # ===================================
-    # index_mask [b, sq, skv]
-    index_mask = torch.full((b, sq, skv), float("-inf"), device=attention_scores.device)
-    index_mask.scatter_(-1, topk_indices, 0)
-    # causal_mask [sq, skv]
-    causal_mask = torch.triu(
-        torch.full((sq, skv), float('-inf'), dtype=torch.float32, device=index_mask.device),
-        diagonal=1,
-    )
-    # [b, sq, skv] + [1, sq, skv] -> [b, sq, skv]
-    index_mask += causal_mask.view(1, sq, skv)
-    # [b, np, sq, skv] + [b, 1, sq, skv] -> [b, np, sq, skv]
-    attention_scores += index_mask.unsqueeze(1)
-    attention_scores = torch.nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32)
+            for s0 in range(0, sq, seq_chunk_size):
+                s1 = min(s0 + seq_chunk_size, sq)
+                s_len = s1 - s0
+                idx_seq_raw = topk_indices[bi, s0:s1]  # [s_len, topk]
+                if idx_seq_raw.dtype != torch.int64 or idx_seq_raw.device != query.device:
+                    idx_seq_raw = idx_seq_raw.to(dtype=torch.int64, device=query.device)
+                valid_seq = idx_seq_raw >= 0
+                idx_seq = idx_seq_raw.clamp(min=0, max=safe_k_max)
+                q_chunk = query_b[bi, h0:h1, s0:s1, :]  # [h_chunk, s_len, hn]
+
+                # These tensors participate in autograd; reusing cached storage can
+                # invalidate saved tensors before backward runs.
+                m = torch.full(
+                    (h_chunk, s_len), float("-inf"), dtype=torch.float32, device=query.device
+                )
+                l = torch.zeros((h_chunk, s_len), dtype=torch.float32, device=query.device)
+                acc = torch.zeros((h_chunk, s_len, hnv), dtype=torch.float32, device=query.device)
+
+                for t0 in range(0, idx_seq.size(-1), topk_chunk_size):
+                    t1 = min(t0 + topk_chunk_size, idx_seq.size(-1))
+                    idx_topk = idx_seq[:, t0:t1]  # [s_len, tk]
+                    valid_t = valid_seq[:, t0:t1]  # [s_len, tk]
+                    flat_idx = idx_topk.unsqueeze(0) + head_offsets  # [h_chunk, s_len, tk]
+                    k_sel = flat_k.index_select(0, flat_idx.reshape(-1)).view(
+                        h_chunk, s_len, -1, hn
+                    )
+                    v_sel = flat_v.index_select(0, flat_idx.reshape(-1)).view(
+                        h_chunk, s_len, -1, hnv
+                    )
+                    logits = (q_chunk.float().unsqueeze(2) * k_sel.float()).sum(
+                        dim=-1
+                    ) * softmax_scale
+
+                    valid_2d, mask_bias = dsa_masking.gather_sparse_topk_validity_and_bias(
+                        idx_topk=idx_topk,
+                        valid_t=valid_t,
+                        bi=bi,
+                        s0=s0,
+                        s1=s1,
+                        row_mask=row_mask,
+                        varlen_starts=varlen_starts,
+                        varlen_ends=varlen_ends,
+                        key_positions=key_positions,
+                        dtype=torch.float32,
+                    )
+                    if mask_bias is not None:
+                        logits = logits + mask_bias.unsqueeze(0)
+                    logits = logits.masked_fill(
+                        ~valid_2d.unsqueeze(0).expand(h_chunk, -1, -1), float("-inf")
+                    )
+                    m_new = torch.maximum(m, logits.max(dim=-1).values)
+                    m_new_for_exp = torch.where(
+                        torch.isfinite(m_new), m_new, torch.zeros_like(m_new)
+                    )
+                    alpha = torch.exp(m - m_new_for_exp)
+                    p = torch.exp(logits - m_new_for_exp.unsqueeze(-1))
+                    acc = acc * alpha.unsqueeze(-1) + torch.einsum(
+                        "hst,hstd->hsd", p, v_sel.float()
+                    )
+                    l = l * alpha + p.sum(dim=-1)
+                    m = m_new
+
+                out_chunk = (acc / l.clamp_min(1e-10).unsqueeze(-1)).to(dtype=value.dtype)
+                output[s0:s1, bi, out_h0:out_h1] = out_chunk.permute(1, 0, 2).reshape(
+                    s_len, h_chunk * hnv
+                )
 
-    # ===================================
-    # Output
-    # ===================================
-    # [skv, b, np, hnv] -> [b, np, skv, hnv] -> [b * np, skv, hnv]
-    value = value.permute(1, 2, 0, 3).reshape(b * np, skv, hnv)
-    # Reshape attention_scores: [b, np, sq, skv] -> [b * np, sq, skv]
-    attention_scores = attention_scores.reshape(b * np, sq, skv)
-    # Compute output: [b * np, sq, hnv]
-    output = torch.bmm(attention_scores.to(value.dtype), value)
-    # Reshape output: [b * np, sq, hnv] -> [b, np, sq, hnv] -> [sq, b, np, hnv]
-    output = output.reshape(b, np, sq, hnv).permute(2, 0, 1, 3).contiguous()
-    # Flatten: [sq, b, np, hnv] -> [sq, b, np * hnv]
-    output = output.reshape(sq, b, np * hnv)
+    if query_was_thd:
+        output = output.squeeze(1)
     return output
 
 
@@ -984,6 +1521,8 @@ class DSAttention(MegatronModule):
         https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/model.py#L491-L597
     """
 
+    consumes_absorbed_v_up_projection = True
+
     def __init__(
         self,
         config: TransformerConfig,
@@ -1011,28 +1550,32 @@ def __init__(
                 k_channels if k_channels is not None else config.kv_channels
             )
         self.softmax_scale = softmax_scale
+        self.cp_comm_type = dsa_layout.normalize_cp_comm_type(cp_comm_type)
 
     def forward(
         self,
         query: torch.Tensor,
         key: torch.Tensor,
-        value: torch.Tensor,
+        value: Optional[torch.Tensor],
         attention_mask: torch.Tensor,
         x: torch.Tensor,
         qr: torch.Tensor,
+        position_ids: Optional[torch.Tensor] = None,
         attn_mask_type: AttnMaskType = None,
         attention_bias: torch.Tensor = None,
         packed_seq_params: PackedSeqParams = None,
+        up_v_weight: Optional[torch.Tensor] = None,
     ):
         """
         Forward pass for Sparse Attention.
 
         Args:
-            query: Query tensor [sq, b, np, hn].
-            key: Key tensor [skv, b, np, hn].
-            value: Value tensor [skv, b, np, hnv].
+            query: Query tensor [sq, b, np, hn] or packed [t, np, hn].
+            key: Key tensor [skv, b, np, hn] or packed [t, np, hn].
+            value: Value tensor [skv, b, np, hnv] or packed [t, np, hnv].
             x: Original hidden states [sq, b, hidden_size].
             qr: Low-rank query representation [sq, b, q_lora_rank].
+            position_ids: Optional position ids [b, sq], used by allgather CP causal masking.
             attention_mask: Attention mask tensor [b, 1, sq, sk].
             attn_mask_type: Type of attention mask.
             attention_bias: Optional attention bias.
@@ -1041,84 +1584,337 @@ def forward(
         Returns:
             output: Output tensor [sq, b, hidden_size]
         """
-        sq, b, np, hn = query.size()
+        query, _ = dsa_layout.ensure_sbhd(query, "query")
+        key, _ = dsa_layout.ensure_sbhd(key, "key")
+        if value is not None:
+            value, _ = dsa_layout.ensure_sbhd(value, "value")
+        if up_v_weight is not None:
+            assert up_v_weight.ndim == 3, "up_v_weight must be [heads, v_head_dim, kv_lora_rank]"
+            up_v_weight = up_v_weight.to(device=query.device, dtype=query.dtype).contiguous()
+            if value is not None:
+                raise RuntimeError(
+                    "DSAttention received up_v_weight with explicit value tensor. "
+                    "For absorbed DSA path, value must be None."
+                )
+
+        latent_v_channels = int(getattr(self.config, "kv_lora_rank", 0) or 0)
+        qk_pos_dim = int(getattr(self.config, "qk_pos_emb_head_dim", 0) or 0)
+        expected_absorbed_dim = latent_v_channels + qk_pos_dim
+        absorbed_mla = (
+            latent_v_channels > 0
+            and expected_absorbed_dim > 0
+            and key.size(2) == 1
+            and query.size(-1) == key.size(-1) == expected_absorbed_dim
+        )
+        if value is None and not absorbed_mla:
+            raise RuntimeError(
+                "DSAttention received value=None but query/key are not in absorbed layout. "
+                f"query_hdim={query.size(-1)}, key_hdim={key.size(-1)}, key_heads={key.size(2)}, "
+                f"expected_absorbed_dim={expected_absorbed_dim}"
+            )
+        if up_v_weight is not None and not absorbed_mla:
+            raise RuntimeError(
+                "DSAttention received up_v_weight but absorbed layout was not detected. "
+                f"query_hdim={query.size(-1)}, key_hdim={key.size(-1)}, key_heads={key.size(2)}, "
+                f"expected_absorbed_dim={expected_absorbed_dim}"
+            )
+
+        sq, b, _, _ = query.size()
+
+        cp_group = getattr(self.indexer.pg_collection, "cp", None)
+        cp_size = cp_group.size() if cp_group is not None else 1
+        cp_rank = cp_group.rank() if cp_group is not None else 0
+        packed_thd = packed_seq_params is not None and packed_seq_params.qkv_format == "thd"
+        packed_query_positions = None
+        kv_reorder_idx = None
+        single_packed_thd_sequence = False
+        if packed_thd and cp_size > 1:
+            cu_seqlens_q, cu_seqlens_kv = dsa_layout.get_packed_qk_cu_seqlens(packed_seq_params)
+            single_packed_thd_sequence = cu_seqlens_q.numel() == 2 and cu_seqlens_kv.numel() == 2
+            packed_query_positions, kv_reorder_idx = (
+                dsa_layout.build_packed_allgather_cp_query_positions_and_key_reorder(
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_kv=cu_seqlens_kv,
+                    cp_size=cp_size,
+                    cp_rank=cp_rank,
+                    device=query.device,
+                    local_output_size=sq,
+                    global_output_size=sq * cp_size,
+                )
+            )
+        elif cp_size > 1:
+            _validate_nonpacked_cp_uniform_length(
+                sq=sq, skv=key.size(0), cp_size=cp_size, cp_group=cp_group, device=query.device
+            )
+            kv_reorder_idx = dsa_layout.build_zigzag_allgather_cp_key_reorder(
+                sq=sq, cp_size=cp_size, device=query.device
+            )
+
+        if cp_size > 1:
+            assert (
+                self.cp_comm_type == "allgather"
+            ), "DSAttention context parallelism currently supports cp_comm_type=allgather only."
+            # For allgather CP, keys/values are expected in full-sequence order.
+            # Gather local-sequence tensors, then undo MCore's zigzag rank order.
+            gathered_cp_key = False
+            gathered_cp_value = False
+            if key.size(0) == sq:
+                key = gather_from_sequence_parallel_region(key, group=cp_group)
+                gathered_cp_key = True
+            if value is not None and value.size(0) == sq:
+                value = gather_from_sequence_parallel_region(value, group=cp_group)
+                gathered_cp_value = True
+            if kv_reorder_idx is not None:
+                if gathered_cp_key:
+                    if key.size(0) != kv_reorder_idx.numel():
+                        raise RuntimeError(
+                            "DSA gathered key length mismatch: "
+                            f"key_seqlen={key.size(0)}, expected={kv_reorder_idx.numel()}"
+                        )
+                    key = key.index_select(0, kv_reorder_idx)
+                if gathered_cp_value:
+                    if value.size(0) != kv_reorder_idx.numel():
+                        raise RuntimeError(
+                            "DSA gathered value length mismatch: "
+                            f"value_seqlen={value.size(0)}, expected={kv_reorder_idx.numel()}"
+                        )
+                    value = value.index_select(0, kv_reorder_idx)
+
         skv = key.size(0)
-        hnv = value.size(3)
 
         # Detach x and qr to prevent gradients of indexer from flowing back to the main model.
         x = x.detach()
         qr = qr.detach()
 
-        # Get a FP32 mask with -inf for masked positions.
-        if attn_mask_type is not None:
-            assert attn_mask_type == AttnMaskType.causal, 'Only causal mask is supported for now'
-            # Generate upper triangular mask with -inf above diagonal, 0 elsewhere
-            # torch.triu with diagonal=1 creates upper triangular matrix (excluding main diagonal)
-            # float_mask [sq, skv]
-            float_mask = torch.triu(
-                torch.full((sq, skv), float('-inf'), dtype=torch.float32, device=x.device),
-                diagonal=1,
-            )
+        indexer_loss_coeff = self.config.dsa_indexer_loss_coeff
+        use_indexer_loss = self.training and torch.is_grad_enabled() and indexer_loss_coeff > 0
+        float_mask, varlen_params = dsa_masking.build_dsattention_forward_mask(
+            sq=sq,
+            skv=skv,
+            b=b,
+            device=x.device,
+            cp_size=cp_size,
+            cp_rank=cp_rank,
+            cp_comm_type=self.cp_comm_type,
+            cp_group=cp_group,
+            attn_mask_type=attn_mask_type,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            packed_seq_params=packed_seq_params,
+            packed_query_positions=packed_query_positions,
+        )
+        if varlen_params is not None:
+            varlen_starts, varlen_ends, key_positions = varlen_params
         else:
-            assert attention_mask.shape == (b, 1, sq, skv), 'attention_mask shape mismatch'
-            # [b, 1, sq, skv] -> [b, sq, skv]
-            mask = attention_mask.squeeze()
-            # float_mask [b, sq, skv]
-            float_mask = torch.zeros_like(mask, dtype=torch.float32).masked_fill(
-                mask, float('-inf')
-            )
-
-        if self.training and torch.is_grad_enabled():
-            # ===================================
-            # Prepare inputs for indexer loss
-            # ===================================
-            q, k, weights = self.indexer.forward_before_topk(x, qr, packed_seq_params)
-            indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0)
+            varlen_starts = varlen_ends = key_positions = None
+        query_valid_rows = dsa_masking.extract_query_valid_rows_from_packed_seq_params(
+            packed_seq_params, b=b, sq=sq, device=query.device
+        )
+        use_fused_kernels = dsa_kernels.use_fused_dsa_kernels(self.config)
+        sparse_indexer_loss = self.config.dsa_indexer_use_sparse_loss
+        use_local_indexer_varlen = (
+            packed_thd
+            and cp_size > 1
+            and single_packed_thd_sequence
+            and attn_mask_type == AttnMaskType.causal
+            and varlen_starts is not None
+            and varlen_ends is not None
+        )
+        indexer_reduce_group = (
+            cp_group if cp_size > 1 and self.config.calculate_per_token_loss else None
+        )
+        indexer_avg_group = (
+            cp_group if cp_size > 1 and not self.config.calculate_per_token_loss else None
+        )
 
-            # ===================================
-            # Attach indexer topk and loss
-            # ===================================
-            # Compute KL divergence loss between indexer scores and true attention scores
-            topk_indices, indexer_loss = FusedDSAIndexerLoss.apply(
+        # ===================================
+        # Prepare indexer inputs / top-k
+        # ===================================
+        q, k, weights = self.indexer.forward_before_topk(x, qr, packed_seq_params)
+        if cp_size > 1 and k.size(0) == sq:
+            k = gather_from_sequence_parallel_region(k, group=cp_group)
+            if kv_reorder_idx is not None:
+                if k.size(0) != kv_reorder_idx.numel():
+                    raise RuntimeError(
+                        "DSA gathered indexer-key length mismatch: "
+                        f"k_seqlen={k.size(0)}, expected={kv_reorder_idx.numel()}"
+                    )
+                k = k.index_select(0, kv_reorder_idx)
+
+        def compute_indexer_loss_with_reference_path():
+            key_for_loss = key.detach()
+            if absorbed_mla and key_for_loss.size(2) == 1 and query.size(2) > 1:
+                key_for_loss = key_for_loss.expand(-1, -1, query.size(2), -1)
+            return FusedDSAIndexerLoss.apply(
                 q,
                 weights,
                 k,
                 query.detach(),
-                key.detach(),
+                key_for_loss,
                 self.softmax_scale,
                 self.indexer.index_topk,
                 indexer_loss_coeff,
                 float_mask,
-                getattr(self.config, "dsa_indexer_use_sparse_loss", False),
+                sparse_indexer_loss,
                 self.indexer.pg_collection,
+                varlen_starts,
+                varlen_ends,
+                key_positions,
+                query_valid_rows,
+                self.config.calculate_per_token_loss,
+                self.config.dsa_indexer_scoring_relu,
             )
-            # Save indexer loss for logging
-            if indexer_loss_coeff > 0:
+
+        fused_output = None
+        if use_fused_kernels:
+            fused_output = dsa_kernels.run_fused_dsa_attention(
+                config=self.config,
+                query=query,
+                key=key,
+                value=value,
+                up_v_weight=up_v_weight,
+                q_indexer=q,
+                k_indexer=k,
+                indexer_weights=weights,
+                indexer_topk=self.indexer.index_topk,
+                softmax_scale=self.softmax_scale,
+                loss_coeff=indexer_loss_coeff,
+                sparse_loss=sparse_indexer_loss,
+                calculate_per_token_loss=self.config.calculate_per_token_loss,
+                absorbed_mla=absorbed_mla,
+                cp_size=cp_size,
+                attn_mask_type=attn_mask_type,
+                packed_seq_params=packed_seq_params,
+                varlen_starts=varlen_starts,
+                varlen_ends=varlen_ends,
+                key_positions=key_positions,
+                query_valid_rows=query_valid_rows,
+                use_relu=self.config.dsa_indexer_scoring_relu,
+                use_local_indexer_varlen=use_local_indexer_varlen,
+            )
+        if fused_output is not None:
+            output, indexer_loss = fused_output
+            if use_indexer_loss:
+                if indexer_loss is None:
+                    raise RuntimeError("Fused DSA attention did not produce a valid indexer loss.")
                 DSAIndexerLossLoggingHelper.save_loss_to_tracker(
                     loss=indexer_loss,
                     layer_number=self.layer_number,
                     num_layers=self.config.num_layers,
+                    reduce_group=indexer_reduce_group,
+                    avg_group=indexer_avg_group,
                 )
+                output = DSAIndexerLossAutoScaler.apply(output, indexer_loss)
+            return _normalize_dsattention_output_rank(output, x.ndim)
+
+        fused_bounds = None
+        if use_fused_kernels:
+            fused_bounds = dsa_masking.build_fused_indexer_varlen_bounds(
+                sq=sq,
+                skv=skv,
+                device=q.device,
+                mask=float_mask,
+                varlen_starts=varlen_starts,
+                varlen_ends=varlen_ends,
+                key_positions=key_positions,
+            )
 
+        topk_indices = None
+        indexer_loss = None
+
+        if use_indexer_loss:
             # ===================================
-            # Run sparse attention kernel
+            # Attach indexer topk and loss
             # ===================================
-            output = unfused_dsa_fn(query, key, value, topk_indices, self.softmax_scale)
+            if sparse_indexer_loss and fused_bounds is not None:
+                starts_i32, ends_i32 = fused_bounds
+                block_size = int(getattr(self, "fused_indexer_block_size", 8192))
+                fused_topk_with_loss = dsa_kernels.run_fused_qk_topk_with_loss(
+                    self.config,
+                    q,
+                    k,
+                    weights,
+                    self.indexer.index_topk,
+                    starts_i32,
+                    ends_i32,
+                    block_size=max(1, block_size),
+                    query=query.detach(),
+                    key=key.detach(),
+                    softmax_scale=self.softmax_scale,
+                    loss_coeff=indexer_loss_coeff,
+                    pg_collection=self.indexer.pg_collection,
+                    query_valid_rows=query_valid_rows,
+                    calculate_per_token_loss=self.config.calculate_per_token_loss,
+                    use_relu=self.config.dsa_indexer_scoring_relu,
+                )
+                if fused_topk_with_loss is not None:
+                    topk_indices, indexer_loss = fused_topk_with_loss
 
-            # Attach loss to output
-            output = DSAIndexerLossAutoScaler.apply(output, indexer_loss)
+            if topk_indices is None or indexer_loss is None:
+                topk_indices, indexer_loss = compute_indexer_loss_with_reference_path()
 
+            # Save indexer loss for logging.
+            if indexer_loss_coeff > 0:
+                DSAIndexerLossLoggingHelper.save_loss_to_tracker(
+                    loss=indexer_loss,
+                    layer_number=self.layer_number,
+                    num_layers=self.config.num_layers,
+                    reduce_group=indexer_reduce_group,
+                    avg_group=indexer_avg_group,
+                )
         else:
             # ===================================
-            # Get index scores and top-k indices
+            # Get top-k indices
             # ===================================
-            _, topk_indices = self.indexer.forward_with_scores(
-                x, qr, mask=float_mask, packed_seq_params=packed_seq_params
-            )
+            if fused_bounds is not None:
+                starts_i32, ends_i32 = fused_bounds
+                block_size = int(getattr(self, "fused_indexer_block_size", 8192))
+                topk_indices = dsa_kernels.run_fused_qk_topk(
+                    self.config,
+                    q,
+                    k,
+                    weights,
+                    self.indexer.index_topk,
+                    starts_i32,
+                    ends_i32,
+                    block_size=max(1, block_size),
+                    use_relu=self.config.dsa_indexer_scoring_relu,
+                )
 
-            # ===================================
-            # Run sparse attention kernel
-            # ===================================
-            output = unfused_dsa_fn(query, key, value, topk_indices, self.softmax_scale)
+            if topk_indices is None:
+                _, topk_indices = fused_qk_topk_naive(
+                    q,
+                    k,
+                    weights,
+                    self.indexer.index_topk,
+                    mask=float_mask,
+                    varlen_starts=varlen_starts,
+                    varlen_ends=varlen_ends,
+                    key_positions=key_positions,
+                    use_relu=self.config.dsa_indexer_scoring_relu,
+                )
 
-        return output
+        # ===================================
+        # Run sparse attention kernel
+        # ===================================
+        output = _run_sparse_attention(
+            absorbed_mla=absorbed_mla,
+            query=query,
+            key=key,
+            value=value,
+            up_v_weight=up_v_weight,
+            topk_indices=topk_indices,
+            softmax_scale=self.softmax_scale,
+            config=self.config,
+            mask=float_mask,
+            varlen_starts=varlen_starts,
+            varlen_ends=varlen_ends,
+            key_positions=key_positions,
+        )
+
+        if use_indexer_loss:
+            if indexer_loss is None:
+                raise RuntimeError("Indexer loss path did not produce a valid loss tensor.")
+            output = DSAIndexerLossAutoScaler.apply(output, indexer_loss)
+
+        return _normalize_dsattention_output_rank(output, x.ndim)
diff --git a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py
new file mode 100644
index 00000000000..5e6b5a37a1e
--- /dev/null
+++ b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py
@@ -0,0 +1,220 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""Backend-neutral hooks for optional fused DeepSeek sparse attention kernels."""
+
+from __future__ import annotations
+
+from importlib import import_module
+from types import ModuleType
+from typing import TYPE_CHECKING, Optional, Tuple
+
+from torch import Tensor
+
+from megatron.core.transformer.enums import AttnBackend, AttnMaskType
+
+if TYPE_CHECKING:
+    from megatron.core.packed_seq_params import PackedSeqParams
+    from megatron.core.process_groups_config import ProcessGroupCollection
+    from megatron.core.transformer.transformer_config import TransformerConfig
+
+_BACKEND_MODULE_NAME_BY_BACKEND = {
+    "tilelang": "megatron.core.transformer.experimental_attention_variant.dsa_tilelang_kernels",
+    "cudnn": "megatron.core.transformer.experimental_attention_variant.dsa_cudnn_kernels",
+}
+_BACKEND: Optional[ModuleType] = None
+_BACKEND_SELECTION: Optional[str] = None
+
+
+def _get_dsa_kernel_backend(config: TransformerConfig) -> str:
+    """Return the configured DSA kernel backend."""
+    backend = config.dsa_kernel_backend
+    if backend != "none" and backend not in _BACKEND_MODULE_NAME_BY_BACKEND:
+        raise ValueError("dsa_kernel_backend must be one of: none, tilelang, cudnn")
+    return backend
+
+
+def _get_backend_module_name(config: TransformerConfig) -> Optional[str]:
+    """Return the optional DSA backend module selected by config."""
+    backend = _get_dsa_kernel_backend(config)
+    if backend == "none":
+        return None
+    return _BACKEND_MODULE_NAME_BY_BACKEND[backend]
+
+
+def _load_backend(config: TransformerConfig) -> Optional[ModuleType]:
+    """Import the configured optional DSA kernel backend."""
+    global _BACKEND, _BACKEND_SELECTION
+    module_name = _get_backend_module_name(config)
+    if module_name is None:
+        _BACKEND = None
+        _BACKEND_SELECTION = None
+        return None
+    if _BACKEND is not None and _BACKEND_SELECTION == module_name:
+        return _BACKEND
+
+    try:
+        _BACKEND = import_module(module_name)
+    except (ImportError, OSError) as exc:
+        raise RuntimeError(f"Failed to import DSA kernel backend {module_name}.") from exc
+    _BACKEND_SELECTION = module_name
+    return _BACKEND
+
+
+def use_fused_dsa_kernels(config: TransformerConfig) -> bool:
+    """Return whether DSA should attempt optional fused kernels before falling back."""
+    backend = config.attention_backend
+    if backend == AttnBackend.unfused or backend == "unfused":
+        return False
+    return _get_dsa_kernel_backend(config) != "none"
+
+
+def run_fused_qk_topk(
+    config: TransformerConfig,
+    q: Tensor,
+    k: Tensor,
+    weights: Tensor,
+    index_topk: int,
+    starts: Tensor,
+    ends: Tensor,
+    block_size: int,
+    use_relu: bool = True,
+) -> Optional[Tensor]:
+    """Optional fused indexer hook for backend-specific implementations."""
+    backend = _load_backend(config)
+    if backend is None:
+        return None
+    fn = getattr(backend, "run_fused_qk_topk", None)
+    if fn is None:
+        return None
+    return fn(q, k, weights, index_topk, starts, ends, block_size, use_relu)
+
+
+def run_fused_qk_topk_with_loss(
+    config: TransformerConfig,
+    q: Tensor,
+    k: Tensor,
+    weights: Tensor,
+    index_topk: int,
+    starts: Tensor,
+    ends: Tensor,
+    block_size: int,
+    query: Tensor,
+    key: Tensor,
+    softmax_scale: float,
+    loss_coeff: float,
+    pg_collection: ProcessGroupCollection,
+    query_valid_rows: Optional[Tensor] = None,
+    calculate_per_token_loss: bool = False,
+    use_relu: bool = True,
+) -> Optional[Tuple[Tensor, Tensor]]:
+    """Optional fused indexer+loss hook for backend-specific implementations."""
+    backend = _load_backend(config)
+    if backend is None:
+        return None
+    fn = getattr(backend, "run_fused_qk_topk_with_loss", None)
+    if fn is None:
+        return None
+    return fn(
+        q=q,
+        k=k,
+        weights=weights,
+        index_topk=index_topk,
+        starts=starts,
+        ends=ends,
+        block_size=block_size,
+        query=query,
+        key=key,
+        softmax_scale=softmax_scale,
+        loss_coeff=loss_coeff,
+        pg_collection=pg_collection,
+        query_valid_rows=query_valid_rows,
+        calculate_per_token_loss=calculate_per_token_loss,
+        use_relu=use_relu,
+    )
+
+
+def run_fused_absorbed_sparse_attention(
+    config: TransformerConfig,
+    query: Tensor,
+    key: Tensor,
+    topk_indices: Tensor,
+    softmax_scale: float,
+    v_channels: int,
+) -> Optional[Tensor]:
+    """Optional fused sparse-attention hook for backend-specific implementations."""
+    backend = _load_backend(config)
+    if backend is None:
+        return None
+    fn = getattr(backend, "run_fused_absorbed_sparse_attention", None)
+    if fn is None:
+        return None
+    return fn(query, key, topk_indices, softmax_scale, v_channels)
+
+
+def run_fused_dsa_attention(
+    *,
+    config: TransformerConfig,
+    query: Tensor,
+    key: Tensor,
+    value: Optional[Tensor],
+    up_v_weight: Optional[Tensor],
+    q_indexer: Tensor,
+    k_indexer: Tensor,
+    indexer_weights: Tensor,
+    indexer_topk: int,
+    softmax_scale: float,
+    loss_coeff: float,
+    sparse_loss: bool,
+    calculate_per_token_loss: bool,
+    absorbed_mla: bool,
+    cp_size: int,
+    attn_mask_type: Optional[AttnMaskType],
+    packed_seq_params: Optional[PackedSeqParams],
+    varlen_starts: Optional[Tensor],
+    varlen_ends: Optional[Tensor],
+    key_positions: Optional[Tensor],
+    query_valid_rows: Optional[Tensor],
+    use_relu: bool,
+    use_local_indexer_varlen: bool = False,
+) -> Optional[Tuple[Tensor, Tensor]]:
+    """Optional full fused DSA hook for backends that fuse indexer and attention together."""
+    backend = _load_backend(config)
+    if backend is None:
+        return None
+    fn = getattr(backend, "run_fused_dsa_attention", None)
+    if fn is None:
+        return None
+    return fn(
+        config=config,
+        query=query,
+        key=key,
+        value=value,
+        up_v_weight=up_v_weight,
+        q_indexer=q_indexer,
+        k_indexer=k_indexer,
+        indexer_weights=indexer_weights,
+        indexer_topk=indexer_topk,
+        softmax_scale=softmax_scale,
+        loss_coeff=loss_coeff,
+        sparse_loss=sparse_loss,
+        calculate_per_token_loss=calculate_per_token_loss,
+        absorbed_mla=absorbed_mla,
+        cp_size=cp_size,
+        attn_mask_type=attn_mask_type,
+        packed_seq_params=packed_seq_params,
+        varlen_starts=varlen_starts,
+        varlen_ends=varlen_ends,
+        key_positions=key_positions,
+        query_valid_rows=query_valid_rows,
+        use_relu=use_relu,
+        use_local_indexer_varlen=use_local_indexer_varlen,
+    )
+
+
+__all__ = [
+    "run_fused_absorbed_sparse_attention",
+    "run_fused_dsa_attention",
+    "run_fused_qk_topk",
+    "run_fused_qk_topk_with_loss",
+    "use_fused_dsa_kernels",
+]
diff --git a/megatron/core/transformer/experimental_attention_variant/dsa_layout.py b/megatron/core/transformer/experimental_attention_variant/dsa_layout.py
new file mode 100644
index 00000000000..fed9889714d
--- /dev/null
+++ b/megatron/core/transformer/experimental_attention_variant/dsa_layout.py
@@ -0,0 +1,283 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""Layout helpers for DeepSeek sparse attention."""
+
+from typing import Optional, Tuple
+
+import torch
+
+from megatron.core.packed_seq_params import PackedSeqParams
+
+__all__ = [
+    "build_packed_allgather_cp_local_positions",
+    "build_packed_allgather_cp_query_positions_and_key_reorder",
+    "build_zigzag_allgather_cp_key_reorder",
+    "build_zigzag_cp_local_positions",
+    "ensure_sbhd",
+    "extract_query_positions_from_position_ids",
+    "get_cp_positions_from_layout",
+    "get_packed_qk_cu_seqlens",
+    "normalize_cp_comm_type",
+]
+
+
+def normalize_cp_comm_type(cp_comm_type: Optional[str]) -> str:
+    """Normalize CP communication type to a canonical lowercase form."""
+    if cp_comm_type is None:
+        return "p2p"
+    return cp_comm_type.replace("_", "").lower()
+
+
+def ensure_sbhd(tensor: torch.Tensor, name: str) -> Tuple[torch.Tensor, bool]:
+    """Ensure tensor is [s, b, h, d], allowing packed [t, h, d] input."""
+    if tensor.ndim == 4:
+        return tensor, False
+    if tensor.ndim == 3:
+        return tensor.unsqueeze(1), True
+    raise ValueError(f"{name} must be 3D ([t,h,d]) or 4D ([s,b,h,d]), got {tensor.ndim}D")
+
+
+def build_zigzag_cp_local_positions(
+    seq_len: int, cp_size: int, cp_rank: int, device: torch.device
+) -> torch.Tensor:
+    """Build this CP rank's token positions under MCore zigzag sequence sharding."""
+    if cp_size <= 1:
+        return torch.arange(seq_len, device=device, dtype=torch.int64)
+    if seq_len % (2 * cp_size) != 0:
+        raise ValueError(
+            "Zigzag CP expects the global sequence length to be divisible by 2 * cp_size, got "
+            f"seq_len={seq_len}, cp_size={cp_size}"
+        )
+
+    chunk_len = seq_len // (2 * cp_size)
+    front_chunk = cp_rank
+    back_chunk = 2 * cp_size - cp_rank - 1
+    return torch.cat(
+        (
+            torch.arange(
+                front_chunk * chunk_len,
+                (front_chunk + 1) * chunk_len,
+                device=device,
+                dtype=torch.int64,
+            ),
+            torch.arange(
+                back_chunk * chunk_len,
+                (back_chunk + 1) * chunk_len,
+                device=device,
+                dtype=torch.int64,
+            ),
+        ),
+        dim=0,
+    )
+
+
+def build_zigzag_allgather_cp_key_reorder(
+    sq: int, cp_size: int, device: torch.device
+) -> torch.Tensor:
+    """Build gathered-KV reorder index for non-packed zigzag allgather CP."""
+    global_seq_len = sq * cp_size
+    gathered_key_positions = torch.cat(
+        [
+            build_zigzag_cp_local_positions(global_seq_len, cp_size, rank, device)
+            for rank in range(cp_size)
+        ],
+        dim=0,
+    )
+    return torch.argsort(gathered_key_positions)
+
+
+def get_cp_positions_from_layout(
+    sq: int,
+    skv: int,
+    cp_size: int,
+    cp_rank: int,
+    cp_comm_type: Optional[str],
+    device: torch.device,
+    cp_group: Optional[torch.distributed.ProcessGroup] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Infer query/key global token positions under CP allgather layout."""
+    if cp_size <= 1:
+        query_pos = torch.arange(sq, device=device, dtype=torch.int64)
+        key_pos = torch.arange(skv, device=device, dtype=torch.int64)
+        return query_pos, key_pos
+
+    if normalize_cp_comm_type(cp_comm_type) != "allgather":
+        raise NotImplementedError(
+            "DSAttention context parallelism currently supports cp_comm_type=allgather only."
+        )
+
+    if skv == sq * cp_size:
+        query_pos = build_zigzag_cp_local_positions(skv, cp_size, cp_rank, device)
+        key_pos = torch.arange(skv, device=device, dtype=torch.int64)
+        return query_pos, key_pos
+
+    # Fallback for callers that pass uneven per-rank lengths. The non-packed MCore
+    # dataloader uses zigzag layout, so the uniform case above is the expected path.
+    query_offset = cp_rank * sq
+    if (
+        cp_group is not None
+        and torch.distributed.is_available()
+        and torch.distributed.is_initialized()
+        and cp_group.size() == cp_size
+    ):
+        local_len = torch.tensor([sq], device=device, dtype=torch.int64)
+        all_lens = [torch.empty_like(local_len) for _ in range(cp_size)]
+        torch.distributed.all_gather(all_lens, local_len, group=cp_group)
+        query_offset = int(torch.stack(all_lens[:cp_rank]).sum().item()) if cp_rank > 0 else 0
+
+    query_pos = torch.arange(sq, device=device, dtype=torch.int64) + query_offset
+    key_pos = torch.arange(skv, device=device, dtype=torch.int64)
+    return query_pos, key_pos
+
+
+def build_packed_allgather_cp_local_positions(
+    cu_seqlens: torch.Tensor,
+    cp_size: int,
+    cp_rank: int,
+    device: torch.device,
+    output_size: Optional[int] = None,
+) -> torch.Tensor:
+    """Build local packed-token positions for one CP rank under zigzag THD sharding.
+
+    This mirrors the packed THD CP layout used by the surrounding training stack:
+    each packed sequence is padded to a multiple of ``2 * cp_size`` and each rank
+    receives the rank-local front chunk followed by the mirrored back chunk.
+    """
+    cu_seqlens_i64 = cu_seqlens.to(device=device, dtype=torch.int64)
+    if cp_size <= 1:
+        if output_size is None:
+            output_size = int(cu_seqlens_i64[-1].item())
+        return torch.arange(output_size, dtype=torch.int64, device=device)
+
+    seq_starts = cu_seqlens_i64[:-1]
+    seq_ends = cu_seqlens_i64[1:]
+    seq_lens = seq_ends - seq_starts
+    nonzero = seq_lens > 0
+    seq_starts = seq_starts[nonzero]
+    seq_ends = seq_ends[nonzero]
+    seq_lens = seq_lens[nonzero]
+    if seq_lens.numel() == 0:
+        return torch.empty(0, dtype=torch.int64, device=device)
+
+    if cu_seqlens_i64.device.type == "cpu":
+        bad_divisible = seq_lens[seq_lens % cp_size != 0]
+        if bad_divisible.numel() > 0:
+            raise ValueError(
+                "Packed DSA CP expects per-sequence padded lengths divisible by cp_size, got "
+                f"seq_len={int(bad_divisible[0].item())}, cp_size={cp_size}"
+            )
+        bad_local = seq_lens[(seq_lens // cp_size) % 2 != 0]
+        if bad_local.numel() > 0:
+            seq_len = int(bad_local[0].item())
+            raise ValueError(
+                "Packed DSA CP expects per-rank packed sequence lengths divisible by 2, got "
+                f"local_seq_len={seq_len // cp_size}, seq_len={seq_len}, cp_size={cp_size}"
+            )
+
+    half_seq_lens = (seq_lens // cp_size) // 2
+    front_starts = seq_starts + cp_rank * half_seq_lens
+    back_starts = seq_ends - (cp_rank + 1) * half_seq_lens
+    segment_starts = torch.stack((front_starts, back_starts), dim=1).reshape(-1)
+    segment_lens = torch.stack((half_seq_lens, half_seq_lens), dim=1).reshape(-1)
+    nonempty_segments = segment_lens > 0
+    segment_starts = segment_starts[nonempty_segments]
+    segment_lens = segment_lens[nonempty_segments]
+
+    if output_size is None:
+        output_size = int(segment_lens.sum().item())
+    if output_size == 0:
+        return torch.empty(0, dtype=torch.int64, device=device)
+
+    segment_ids = torch.repeat_interleave(
+        torch.arange(segment_lens.numel(), dtype=torch.int64, device=device),
+        segment_lens,
+        output_size=output_size,
+    )
+    segment_offsets = torch.arange(output_size, dtype=torch.int64, device=device)
+    segment_offsets -= torch.repeat_interleave(
+        torch.cumsum(segment_lens, dim=0) - segment_lens, segment_lens, output_size=output_size
+    )
+    return segment_starts.index_select(0, segment_ids) + segment_offsets
+
+
+def build_packed_allgather_cp_query_positions_and_key_reorder(
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_kv: torch.Tensor,
+    cp_size: int,
+    cp_rank: int,
+    device: torch.device,
+    local_output_size: Optional[int] = None,
+    global_output_size: Optional[int] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Build packed-query positions and gathered-KV reorder index for allgather CP.
+
+    Queries stay in the local zigzag THD order for ``cp_rank``. Keys/values are
+    manually all-gathered rank-by-rank, so their gathered tensor order is:
+    rank0-local-packed, rank1-local-packed, ..., rank{cp_size-1}-local-packed.
+    This helper returns the permutation that restores those gathered KV tensors
+    to global packed order, matching the Slime GLM5 implementation semantics.
+    """
+    query_positions = build_packed_allgather_cp_local_positions(
+        cu_seqlens_q, cp_size, cp_rank, device, output_size=local_output_size
+    )
+    gathered_key_positions = [
+        build_packed_allgather_cp_local_positions(
+            cu_seqlens_kv, cp_size, rank, device, output_size=local_output_size
+        )
+        for rank in range(cp_size)
+    ]
+    gathered_key_positions = torch.cat(gathered_key_positions, dim=0)
+    key_reorder_idx = torch.argsort(gathered_key_positions)
+    if global_output_size is not None and key_reorder_idx.numel() != global_output_size:
+        raise RuntimeError(
+            f"Packed DSA CP key reorder length mismatch: got {key_reorder_idx.numel()}, "
+            f"expected {global_output_size}"
+        )
+    return query_positions, key_reorder_idx
+
+
+def extract_query_positions_from_position_ids(
+    position_ids: Optional[torch.Tensor], sq: int, device: torch.device
+) -> Optional[torch.Tensor]:
+    """Extract per-rank query positions from position_ids if compatible."""
+    if position_ids is None:
+        return None
+    if position_ids.ndim == 2:
+        if position_ids.size(0) > 1:
+            assert torch.equal(
+                position_ids[0], position_ids[-1]
+            ), "Allgather-CP DSA expects identical position_ids across batch"
+        query_pos = position_ids[0]
+    elif position_ids.ndim == 1:
+        query_pos = position_ids
+    else:
+        raise ValueError(f"position_ids should be 1D or 2D tensor, got {position_ids.ndim}D.")
+
+    if query_pos.numel() != sq:
+        return None
+    return query_pos.to(device=device, dtype=torch.int64)
+
+
+def get_packed_qk_cu_seqlens(
+    packed_seq_params: PackedSeqParams,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Select packed cu_seqlens for query and key/value streams."""
+    cu_seqlens_q = (
+        packed_seq_params.cu_seqlens_q_padded
+        if packed_seq_params.cu_seqlens_q_padded is not None
+        else packed_seq_params.cu_seqlens_q
+    )
+    cu_seqlens = (
+        packed_seq_params.cu_seqlens_kv_padded
+        if packed_seq_params.cu_seqlens_kv_padded is not None
+        else packed_seq_params.cu_seqlens_kv
+    )
+    cu_seqlens_kv = cu_seqlens
+
+    if cu_seqlens_q is None and cu_seqlens_kv is None:
+        raise ValueError("Packed sequence parameters must provide cu_seqlens for DSA masking.")
+    if cu_seqlens_q is None:
+        cu_seqlens_q = cu_seqlens_kv
+    if cu_seqlens_kv is None:
+        cu_seqlens_kv = cu_seqlens_q
+    return cu_seqlens_q, cu_seqlens_kv
diff --git a/megatron/core/transformer/experimental_attention_variant/dsa_masking.py b/megatron/core/transformer/experimental_attention_variant/dsa_masking.py
new file mode 100644
index 00000000000..5fdda3c9a4d
--- /dev/null
+++ b/megatron/core/transformer/experimental_attention_variant/dsa_masking.py
@@ -0,0 +1,505 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""Masking helpers for DeepSeek sparse attention."""
+
+from typing import Optional, Tuple
+
+import torch
+
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.experimental_attention_variant import dsa_layout
+
+__all__ = [
+    "apply_sparse_validity_to_index_mask",
+    "apply_starts_ends_mask_to_scores",
+    "build_causal_mask_from_positions",
+    "build_dsattention_forward_mask",
+    "build_fused_indexer_varlen_bounds",
+    "build_valid_mask_from_starts_ends",
+    "extract_query_valid_rows_from_packed_seq_params",
+    "gather_sparse_topk_validity_and_bias",
+    "generate_varlen_mask_params",
+    "generate_varlen_mask_params_for_positions",
+    "masked_softmax",
+    "masked_softmax_inplace",
+    "normalize_query_valid_rows",
+    "normalize_varlen_bounds",
+    "prepare_additive_mask",
+    "prepare_sparse_mask_context",
+    "scatter_topk_into_index_mask",
+]
+
+
+def build_causal_mask_from_positions(
+    query_pos: torch.Tensor, key_pos: torch.Tensor
+) -> torch.Tensor:
+    """Build a causal mask from explicit query/key global positions."""
+    assert query_pos.dtype in (torch.int32, torch.int64), "query_pos must be integer tensor"
+    assert key_pos.dtype in (torch.int32, torch.int64), "key_pos must be integer tensor"
+    assert query_pos.device == key_pos.device, "query_pos and key_pos must be on the same device"
+
+    # mask[q, k] = -inf if key_pos[k] > query_pos[q], else 0.
+    invalid = key_pos.unsqueeze(0) > query_pos.unsqueeze(-1)
+    mask = torch.zeros(
+        (query_pos.numel(), key_pos.numel()), dtype=torch.float32, device=query_pos.device
+    )
+    mask.masked_fill_(invalid, float("-inf"))
+    return mask
+
+
+def generate_varlen_mask_params(cu_seqlens: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Generate row-wise [start, end) key bounds for packed causal masking."""
+    assert cu_seqlens.ndim == 1 and cu_seqlens.numel() >= 2, "invalid cu_seqlens"
+    cu_seqlens = cu_seqlens.to(dtype=torch.int64)
+    seq_len = int(cu_seqlens[-1].item())
+    q_indices = torch.arange(seq_len, dtype=torch.int64, device=cu_seqlens.device)
+    seq_indices = torch.searchsorted(cu_seqlens, q_indices, right=True) - 1
+    starts = cu_seqlens[seq_indices]
+    ends = q_indices + 1
+    return starts, ends
+
+
+def generate_varlen_mask_params_for_positions(
+    cu_seqlens: torch.Tensor, query_positions: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Generate packed causal bounds only for the requested query positions."""
+    assert cu_seqlens.ndim == 1 and cu_seqlens.numel() >= 2, "invalid cu_seqlens"
+    assert query_positions.dtype in (torch.int32, torch.int64), "query_positions must be integer"
+    cu_seqlens = cu_seqlens.to(device=query_positions.device, dtype=torch.int64)
+    query_positions = query_positions.to(dtype=torch.int64)
+    seq_indices = torch.searchsorted(cu_seqlens[1:], query_positions, right=True)
+    starts = cu_seqlens[seq_indices]
+    ends = query_positions + 1
+    return starts, ends
+
+
+def build_valid_mask_from_starts_ends(
+    starts: torch.Tensor, ends: torch.Tensor, key_positions: torch.Tensor
+) -> torch.Tensor:
+    """Build boolean validity mask [sq, sk] from row-wise [start, end) bounds."""
+    assert starts.ndim == ends.ndim == 1, "starts/ends must be 1D"
+    assert starts.shape == ends.shape, "starts/ends shape mismatch"
+    assert key_positions.ndim == 1, "key_positions must be 1D"
+    assert starts.device == ends.device == key_positions.device, "device mismatch"
+    assert starts.dtype in (torch.int32, torch.int64), "starts must be int tensor"
+    assert ends.dtype in (torch.int32, torch.int64), "ends must be int tensor"
+    assert key_positions.dtype in (torch.int32, torch.int64), "key_positions must be int tensor"
+    key_positions = key_positions.to(dtype=torch.int64)
+    starts = starts.to(dtype=torch.int64)
+    ends = ends.to(dtype=torch.int64)
+    return (key_positions.unsqueeze(0) >= starts.unsqueeze(-1)) & (
+        key_positions.unsqueeze(0) < ends.unsqueeze(-1)
+    )
+
+
+def apply_starts_ends_mask_to_scores(
+    scores: torch.Tensor, starts: torch.Tensor, ends: torch.Tensor, key_positions: torch.Tensor
+) -> torch.Tensor:
+    """Apply varlen starts/ends mask to score tensor.
+
+    Supports scores with shape [b, sq, sk] or [b, np, sq, sk].
+    """
+    valid = build_valid_mask_from_starts_ends(starts, ends, key_positions)
+    if scores.ndim == 3:
+        return scores.masked_fill(~valid.unsqueeze(0), float("-inf"))
+    if scores.ndim == 4:
+        return scores.masked_fill(~valid.unsqueeze(0).unsqueeze(0), float("-inf"))
+    raise ValueError(f"Unsupported scores ndim={scores.ndim}, expected 3 or 4.")
+
+
+def normalize_varlen_bounds(
+    *,
+    mask: Optional[torch.Tensor],
+    varlen_starts: Optional[torch.Tensor],
+    varlen_ends: Optional[torch.Tensor],
+    key_positions: Optional[torch.Tensor],
+    sk: int,
+    device: torch.device,
+) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
+    """Validate mask/varlen exclusivity and normalize varlen bounds to int64 tensors."""
+    if mask is not None and varlen_starts is not None:
+        raise ValueError("mask and varlen_starts are mutually exclusive")
+    if varlen_starts is None:
+        return None, None, None
+    if varlen_ends is None:
+        raise ValueError("varlen_ends is required when varlen_starts is provided")
+
+    varlen_starts_i64 = varlen_starts.to(device=device, dtype=torch.int64)
+    varlen_ends_i64 = varlen_ends.to(device=device, dtype=torch.int64)
+    if key_positions is None:
+        key_positions_i64 = torch.arange(sk, dtype=torch.int64, device=device)
+    else:
+        key_positions_i64 = key_positions.to(device=device, dtype=torch.int64)
+    return varlen_starts_i64, varlen_ends_i64, key_positions_i64
+
+
+def _build_default_causal_mask(sq: int, sk: int, device: torch.device) -> torch.Tensor:
+    """Build standard upper-triangular additive causal mask."""
+    return torch.triu(
+        torch.full((sq, sk), float("-inf"), dtype=torch.float32, device=device), diagonal=1
+    )
+
+
+def prepare_additive_mask(
+    mask: Optional[torch.Tensor], *, sq: int, sk: int, b: int, device: torch.device
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Validate/build additive mask and return useful broadcasted views.
+
+    Returns:
+        score_mask: [sq, sk] or [b, sq, sk]
+        attn_score_mask: [1, 1, sq, sk] or [b, 1, sq, sk]
+        index_score_mask: [1, sq, sk] or [b, sq, sk]
+        valid_mask: [b, sq, sk] bool, True means finite (not masked)
+    """
+    if mask is None:
+        score_mask = _build_default_causal_mask(sq, sk, device=device)
+    else:
+        assert mask.dtype == torch.float32, "mask dtype must be float32"
+        assert mask.device == device, "mask device mismatch"
+        assert mask.ndim in (2, 3), "mask must be 2D or 3D"
+        if mask.ndim == 2:
+            assert mask.shape == (sq, sk), "mask shape mismatch"
+        else:
+            assert mask.shape == (b, sq, sk), "mask shape mismatch"
+        score_mask = mask
+
+    if score_mask.ndim == 2:
+        attn_score_mask = score_mask.view(1, 1, sq, sk)
+        index_score_mask = score_mask.unsqueeze(0)
+        valid_mask = torch.isfinite(score_mask).unsqueeze(0).expand(b, sq, sk)
+    else:
+        attn_score_mask = score_mask.view(b, 1, sq, sk)
+        index_score_mask = score_mask
+        valid_mask = torch.isfinite(score_mask)
+    return score_mask, attn_score_mask, index_score_mask, valid_mask
+
+
+def prepare_sparse_mask_context(
+    *,
+    mask: Optional[torch.Tensor],
+    varlen_starts: Optional[torch.Tensor],
+    varlen_ends: Optional[torch.Tensor],
+    key_positions: Optional[torch.Tensor],
+    sq: int,
+    sk: int,
+    b: int,
+    device: torch.device,
+) -> Tuple[
+    Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]
+]:
+    """Prepare shared sparse-mask context for unfused attention paths."""
+    varlen_starts_i64, varlen_ends_i64, key_positions_i64 = normalize_varlen_bounds(
+        mask=mask,
+        varlen_starts=varlen_starts,
+        varlen_ends=varlen_ends,
+        key_positions=key_positions,
+        sk=sk,
+        device=device,
+    )
+    if varlen_starts_i64 is not None:
+        return None, varlen_starts_i64, varlen_ends_i64, key_positions_i64
+
+    _, _, index_score_mask, _ = prepare_additive_mask(mask, sq=sq, sk=sk, b=b, device=device)
+    return index_score_mask, None, None, None
+
+
+def apply_sparse_validity_to_index_mask(
+    index_mask: torch.Tensor,
+    *,
+    row_mask: Optional[torch.Tensor],
+    varlen_starts: Optional[torch.Tensor],
+    varlen_ends: Optional[torch.Tensor],
+    key_positions: Optional[torch.Tensor],
+) -> torch.Tensor:
+    """Apply either varlen or additive mask validity constraints to index_mask."""
+    if varlen_starts is not None:
+        varlen_starts, varlen_ends, key_positions = normalize_varlen_bounds(
+            mask=None,
+            varlen_starts=varlen_starts,
+            varlen_ends=varlen_ends,
+            key_positions=key_positions,
+            sk=index_mask.size(-1),
+            device=index_mask.device,
+        )
+        valid_mask = build_valid_mask_from_starts_ends(
+            varlen_starts, varlen_ends, key_positions
+        ).unsqueeze(0)
+        return index_mask.masked_fill(~valid_mask, float("-inf"))
+
+    if row_mask is None:
+        raise ValueError("row_mask is required when varlen_starts is None")
+    return index_mask + row_mask
+
+
+def gather_sparse_topk_validity_and_bias(
+    *,
+    idx_topk: torch.Tensor,
+    valid_t: torch.Tensor,
+    bi: int,
+    s0: int,
+    s1: int,
+    row_mask: Optional[torch.Tensor],
+    varlen_starts: Optional[torch.Tensor],
+    varlen_ends: Optional[torch.Tensor],
+    key_positions: Optional[torch.Tensor],
+    dtype: torch.dtype,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """Gather top-k validity mask and optional additive bias for one [s_chunk, topk] block."""
+    if varlen_starts is not None:
+        if varlen_ends is None:
+            raise ValueError("varlen_ends is required when varlen_starts is provided")
+        if key_positions is None:
+            raise ValueError("key_positions is required when varlen_starts is provided")
+        key_pos_sel = key_positions.index_select(0, idx_topk.reshape(-1)).view_as(idx_topk)
+        valid_varlen = (key_pos_sel >= varlen_starts[s0:s1].unsqueeze(-1)) & (
+            key_pos_sel < varlen_ends[s0:s1].unsqueeze(-1)
+        )
+        return valid_t & valid_varlen, None
+
+    if row_mask is None:
+        raise ValueError("row_mask is required when varlen_starts is None")
+    mask_src = row_mask[0, s0:s1, :] if row_mask.size(0) == 1 else row_mask[bi, s0:s1, :]
+    mask_bias = mask_src.gather(-1, idx_topk).to(dtype=dtype)
+    return valid_t & torch.isfinite(mask_bias), mask_bias
+
+
+def scatter_topk_into_index_mask(
+    index_mask: torch.Tensor, topk_indices: torch.Tensor, *, seq_chunk_size: int = 256
+) -> None:
+    """Scatter top-k supports into index_mask using chunk-wise int64 casts."""
+    b, sq, _ = index_mask.shape
+    assert topk_indices.ndim == 3, "topk_indices must be [b, sq, topk]"
+    assert topk_indices.shape[:2] == (b, sq), "topk_indices shape mismatch"
+    device = index_mask.device
+    seq_chunk_size = max(1, int(seq_chunk_size))
+
+    for s0 in range(0, sq, seq_chunk_size):
+        s1 = min(s0 + seq_chunk_size, sq)
+        idx_chunk = topk_indices[:, s0:s1]
+        if idx_chunk.dtype != torch.int64 or idx_chunk.device != device:
+            idx_chunk = idx_chunk.to(dtype=torch.int64, device=device)
+        if torch.any(idx_chunk < 0):
+            valid_topk = idx_chunk >= 0
+            if valid_topk.any():
+                b_idx, q_rel_idx, t_idx = torch.where(valid_topk)
+                q_idx = q_rel_idx + s0
+                k_idx = idx_chunk[b_idx, q_rel_idx, t_idx]
+                index_mask[b_idx, q_idx, k_idx] = 0.0
+        else:
+            index_mask[:, s0:s1].scatter_(-1, idx_chunk, 0.0)
+
+
+def masked_softmax_inplace(
+    logits: torch.Tensor, valid_mask: torch.Tensor, *, dim: int = -1, eps: float = 1e-10
+) -> torch.Tensor:
+    """Convert logits to probabilities in place while zeroing invalid entries."""
+    if not logits.is_floating_point():
+        raise TypeError("masked_softmax_inplace expects a floating-point tensor")
+    if logits.shape != valid_mask.shape:
+        raise ValueError("logits and valid_mask must have the same shape")
+
+    logits.masked_fill_(~valid_mask, torch.finfo(logits.dtype).min)
+    row_has_valid = valid_mask.any(dim=dim, keepdim=True)
+    row_max = logits.max(dim=dim, keepdim=True).values
+    row_max = torch.where(row_has_valid, row_max, torch.zeros_like(row_max))
+
+    logits.sub_(row_max)
+    logits.exp_()
+    logits.masked_fill_(~valid_mask, 0.0)
+    logits.div_(logits.sum(dim=dim, keepdim=True).clamp_min(eps))
+    logits.masked_fill_(~valid_mask, 0.0)
+    return logits
+
+
+def masked_softmax(
+    logits: torch.Tensor, valid_mask: torch.Tensor, *, dim: int = -1, eps: float = 1e-10
+) -> torch.Tensor:
+    """Convert logits to probabilities while zeroing invalid entries."""
+    if not logits.is_floating_point():
+        raise TypeError("masked_softmax expects a floating-point tensor")
+    if logits.shape != valid_mask.shape:
+        raise ValueError("logits and valid_mask must have the same shape")
+
+    masked_logits = logits.masked_fill(~valid_mask, torch.finfo(logits.dtype).min)
+    row_has_valid = valid_mask.any(dim=dim, keepdim=True)
+    row_max = masked_logits.max(dim=dim, keepdim=True).values
+    row_max = torch.where(row_has_valid, row_max, torch.zeros_like(row_max))
+
+    probs = torch.exp(masked_logits - row_max)
+    probs = probs.masked_fill(~valid_mask, 0.0)
+    probs = probs / probs.sum(dim=dim, keepdim=True).clamp_min(eps)
+    return probs.masked_fill(~valid_mask, 0.0)
+
+
+def normalize_query_valid_rows(
+    query_valid_rows: Optional[torch.Tensor], *, b: int, sq: int, device: torch.device
+) -> Optional[torch.Tensor]:
+    """Normalize optional query-row validity mask to shape [b, sq]."""
+    if query_valid_rows is None:
+        return None
+    query_valid_rows = query_valid_rows.to(device=device, dtype=torch.bool)
+    if query_valid_rows.ndim == 1:
+        if query_valid_rows.numel() != sq:
+            raise ValueError(
+                f"query_valid_rows length mismatch: expected {sq}, got {query_valid_rows.numel()}"
+            )
+        return query_valid_rows.unsqueeze(0).expand(b, sq)
+    if query_valid_rows.ndim == 2:
+        if query_valid_rows.shape == (1, sq):
+            return query_valid_rows.expand(b, sq)
+        if query_valid_rows.shape != (b, sq):
+            expected_shape = (b, sq)
+            raise ValueError(
+                f"query_valid_rows shape mismatch: expected {expected_shape}, "
+                f"got {tuple(query_valid_rows.shape)}"
+            )
+        return query_valid_rows
+    raise ValueError(f"query_valid_rows should be 1D or 2D tensor, got {query_valid_rows.ndim}D.")
+
+
+def extract_query_valid_rows_from_packed_seq_params(
+    packed_seq_params: Optional[PackedSeqParams], *, b: int, sq: int, device: torch.device
+) -> Optional[torch.Tensor]:
+    """Extract optional real-token query-row mask from packed sequence metadata."""
+    if packed_seq_params is None:
+        return None
+    query_valid_rows = getattr(packed_seq_params, "real_token_mask_q", None)
+    if query_valid_rows is None:
+        return None
+    return normalize_query_valid_rows(query_valid_rows, b=b, sq=sq, device=device)
+
+
+def build_dsattention_forward_mask(
+    *,
+    sq: int,
+    skv: int,
+    b: int,
+    device: torch.device,
+    cp_size: int,
+    cp_rank: int,
+    cp_comm_type: str,
+    cp_group: Optional[torch.distributed.ProcessGroup],
+    attn_mask_type: Optional[AttnMaskType],
+    attention_mask: Optional[torch.Tensor],
+    position_ids: Optional[torch.Tensor],
+    packed_seq_params: Optional[PackedSeqParams],
+    packed_query_positions: Optional[torch.Tensor] = None,
+) -> Tuple[Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]]:
+    """Build DSAttention mask.
+
+    Returns:
+        float_mask: Optional additive mask [sq, skv] or [b, sq, skv].
+        varlen_params: Optional (starts, ends, key_positions), each int64 tensor.
+    """
+    packed_thd = packed_seq_params is not None and packed_seq_params.qkv_format == "thd"
+    if attn_mask_type is not None:
+        assert attn_mask_type == AttnMaskType.causal, "Only causal mask is supported for now"
+        if packed_thd:
+            cu_seqlens_q, _ = dsa_layout.get_packed_qk_cu_seqlens(packed_seq_params)
+            cu_seqlens_q = cu_seqlens_q.to(device=device, dtype=torch.int64)
+            if cp_size > 1:
+                if packed_query_positions is not None:
+                    query_idx = packed_query_positions.to(device=device, dtype=torch.int64)
+                    key_idx = torch.arange(skv, dtype=torch.int64, device=device)
+                else:
+                    query_idx, key_idx = dsa_layout.get_cp_positions_from_layout(
+                        sq=sq,
+                        skv=skv,
+                        cp_size=cp_size,
+                        cp_rank=cp_rank,
+                        cp_comm_type=cp_comm_type,
+                        device=device,
+                        cp_group=cp_group,
+                    )
+            else:
+                query_idx = torch.arange(sq, dtype=torch.int64, device=device)
+                key_idx = torch.arange(skv, dtype=torch.int64, device=device)
+            varlen_starts, varlen_ends = generate_varlen_mask_params_for_positions(
+                cu_seqlens_q, query_idx
+            )
+            return None, (varlen_starts, varlen_ends, key_idx)
+
+        if cp_size > 1:
+            query_pos = dsa_layout.extract_query_positions_from_position_ids(
+                position_ids, sq, device
+            )
+            if query_pos is None:
+                query_pos, key_pos = dsa_layout.get_cp_positions_from_layout(
+                    sq=sq,
+                    skv=skv,
+                    cp_size=cp_size,
+                    cp_rank=cp_rank,
+                    cp_comm_type=cp_comm_type,
+                    device=device,
+                    cp_group=cp_group,
+                )
+            else:
+                key_pos = torch.arange(skv, dtype=torch.int64, device=device)
+            return build_causal_mask_from_positions(query_pos, key_pos), None
+
+        return _build_default_causal_mask(sq, skv, device=device), None
+
+    assert attention_mask is not None, "attention_mask is required when attn_mask_type is None"
+    assert attention_mask.shape == (b, 1, sq, skv), "attention_mask shape mismatch"
+    mask = attention_mask[:, 0, :, :]
+    float_mask = torch.zeros_like(mask, dtype=torch.float32).masked_fill(mask, float("-inf"))
+    return float_mask, None
+
+
+def build_fused_indexer_varlen_bounds(
+    *,
+    sq: int,
+    skv: int,
+    device: torch.device,
+    mask: Optional[torch.Tensor],
+    varlen_starts: Optional[torch.Tensor],
+    varlen_ends: Optional[torch.Tensor],
+    key_positions: Optional[torch.Tensor],
+) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
+    """Build row-wise contiguous [start, end) key bounds for optional fused indexer kernels."""
+    varlen_starts, varlen_ends, key_positions = normalize_varlen_bounds(
+        mask=mask,
+        varlen_starts=varlen_starts,
+        varlen_ends=varlen_ends,
+        key_positions=key_positions,
+        sk=skv,
+        device=device,
+    )
+    if varlen_starts is not None:
+        expected_key_pos = torch.arange(skv, dtype=torch.int64, device=device)
+        if not torch.equal(key_positions, expected_key_pos):
+            return None
+        return (
+            varlen_starts.to(dtype=torch.int32, device=device),
+            varlen_ends.to(dtype=torch.int32, device=device),
+        )
+
+    if mask is None:
+        ends = torch.arange(1, sq + 1, dtype=torch.int64, device=device).clamp_max(skv)
+        starts = torch.zeros_like(ends)
+        return starts.to(dtype=torch.int32), ends.to(dtype=torch.int32)
+
+    if mask.ndim == 3:
+        # Fused indexers generally use one shared bounds schedule. For batched masks, only
+        # enable a fused path when all batch masks are identical.
+        if mask.size(0) > 1:
+            ref_mask = mask[0]
+            for bi in range(1, mask.size(0)):
+                if not torch.equal(mask[bi], ref_mask):
+                    return None
+        row_mask = mask[0]
+    else:
+        row_mask = mask
+    if row_mask.ndim != 2 or row_mask.shape != (sq, skv):
+        return None
+
+    finite = torch.isfinite(row_mask)
+    ends = finite.sum(dim=-1, dtype=torch.int64)
+    key_ids = torch.arange(skv, dtype=torch.int64, device=device).unsqueeze(0)
+    expected = key_ids < ends.unsqueeze(-1)
+    if not torch.equal(finite, expected):
+        return None
+
+    starts = torch.zeros_like(ends)
+    return starts.to(dtype=torch.int32), ends.to(dtype=torch.int32)
diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py
index 5db3154f552..a35505e4548 100644
--- a/megatron/core/transformer/multi_latent_attention.py
+++ b/megatron/core/transformer/multi_latent_attention.py
@@ -15,7 +15,6 @@
 except ImportError:
     HAVE_EINOPS = False
 
-
 from megatron.core import tensor_parallel
 from megatron.core.dist_checkpointing.mapping import ShardedObject
 from megatron.core.extensions.transformer_engine import HAVE_TE
@@ -382,12 +381,6 @@ def forward(
             )
         else:
             if inference_context is None or inference_context.is_static_batching():
-                extra_kwargs = {}
-                if self.config.experimental_attention_variant == "dsa":
-                    # For dsa we need to pass in the original hidden states and the compressed
-                    # query representation.
-                    extra_kwargs["x"] = hidden_states
-                    extra_kwargs["qr"] = q_compressed
                 with off_interface(
                     self.offload_core_attention and self.training, query, "core_attn"
                 ) as query:
@@ -398,7 +391,6 @@ def forward(
                         attention_mask,
                         packed_seq_params=packed_seq_params,
                         attn_mask_type=attn_mask_type,
-                        **extra_kwargs,
                     )
             elif self.cache_mla_latents:
                 value, need_v_pad, orig_v_dim, padded_v_dim = _prepare_mla_core_attention_value(
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index ab84abd3a17..36c0adbc861 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -41,6 +41,8 @@
 
 logger = logging.getLogger(__name__)
 
+_VALID_DSA_KERNEL_BACKENDS = ("none", "tilelang", "cudnn")
+
 try:
     from packaging.version import Version as PkgVersion
 
@@ -49,6 +51,64 @@
     HAVE_PACKAGING = False
 
 
+def _missing_tilelang_dsa_kernel_dependencies() -> List[str]:
+    """Return missing TileLang DSA kernel dependencies."""
+    try:
+        from megatron.core.transformer.experimental_attention_variant.ops import tilelang_dsa
+    except (ImportError, OSError):
+        return ["TileLang DSA kernels"]
+
+    missing = []
+    if tilelang_dsa.lighting_indexer is None:
+        missing.append("TileLang DSA indexer")
+    if tilelang_dsa.SparseMLA is None:
+        missing.append("TileLang SparseMLA")
+    return missing
+
+
+def _missing_cudnn_dsa_kernel_dependencies() -> List[str]:
+    """Return missing cuDNN DSA kernel dependencies."""
+    missing = []
+    try:
+        from flash_mla import flash_mla_sparse_fwd  # noqa: F401
+    except ImportError:
+        missing.append("flash_mla")
+
+    try:
+        from cudnn import DSA  # noqa: F401
+    except ImportError:
+        missing.append("cudnn-frontend DSA (nvidia-cudnn-frontend[cutedsl])")
+    return missing
+
+
+def _validate_dsa_kernel_backend_dependencies(dsa_kernel_backend: str) -> None:
+    """Validate optional fused DSA kernel backend dependencies."""
+    if dsa_kernel_backend not in _VALID_DSA_KERNEL_BACKENDS:
+        raise ValueError(
+            "dsa_kernel_backend must be one of: " f"{', '.join(_VALID_DSA_KERNEL_BACKENDS)}."
+        )
+    if dsa_kernel_backend == "none":
+        return
+    if not torch.cuda.is_available():
+        raise ValueError(
+            f"dsa_kernel_backend={dsa_kernel_backend} requires a CUDA device, "
+            "but none is available."
+        )
+
+    missing = []
+    if dsa_kernel_backend == "tilelang":
+        missing = _missing_tilelang_dsa_kernel_dependencies()
+    elif dsa_kernel_backend == "cudnn":
+        missing = _missing_cudnn_dsa_kernel_dependencies()
+
+    if missing:
+        raise ValueError(
+            f"dsa_kernel_backend={dsa_kernel_backend} requires fused DSA kernels, "
+            f"but the following packages are not available: {', '.join(missing)}. "
+            "Install them or set dsa_kernel_backend=none to use the PyTorch fallback."
+        )
+
+
 @dataclass
 class TransformerConfig(ModelParallelConfig):
     """Configuration object for megatron-core transformers.
@@ -283,6 +343,9 @@ class TransformerConfig(ModelParallelConfig):
     experimental_attention_variant: Optional[Literal['gated_delta_net', 'dsa']] = None
     """Type of attention variant to use. Currently support gated_delta_net and dsa."""
 
+    experimental_attention_variant_loss_scale_func: Optional[Callable[[torch.Tensor], None]] = None
+    """Optional hook for experimental attention variants to receive the main loss scale."""
+
     ####################
     # DSA
     ####################
@@ -302,6 +365,26 @@ class TransformerConfig(ModelParallelConfig):
     """Whether to use sparse DSA indexer loss. If True, the indexer loss will be computed using the
     top-k indices."""
 
+    dsa_kernel_backend: Literal["none", "tilelang", "cudnn"] = "none"
+    """Optional fused DSA kernel backend.
+    ``none`` disables fused DSA kernels. Explicit ``tilelang`` or ``cudnn`` enables only that
+    backend. Unsupported DSA layouts continue to use the PyTorch fallback."""
+
+    dsa_indexer_rope_interleaved: bool = False
+    """Whether DSA indexer RoPE should use MLA-style interleaving."""
+
+    dsa_indexer_rotate_activation: bool = True
+    """Whether DSA indexer should apply Hadamard rotate_activation to q/k before scoring."""
+
+    dsa_indexer_scoring_relu: bool = True
+    """Whether DSA indexer should apply ReLU to q@k^T scores before weighting."""
+
+    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."""
+
     ####################
     # linear attention
     ####################
@@ -1260,7 +1343,7 @@ def __post_init__(self):
                 f"({self.tensor_model_parallel_size=} * {self.context_parallel_size=})."
             )
         elif self.experimental_attention_variant == "dsa":
-            pass
+            _validate_dsa_kernel_backend_dependencies(self.dsa_kernel_backend)
 
         if self.fp8:
             # cannot support first last layer bf16 with delayed scaling
@@ -2557,10 +2640,21 @@ def _scope_to_str(s):
             assert not self.use_kitchen
 
         if self.experimental_attention_variant == "dsa":
-            assert (
-                self.context_parallel_size == 1
-            ), "Currently context parallelism is not supported by DSAttention!"
             assert not self.apply_rope_fusion, "RoPE fusion is not supported for DSAttention"
+            if self.context_parallel_size > 1:
+                cp_comm_types = (
+                    self.cp_comm_type
+                    if isinstance(self.cp_comm_type, list)
+                    else [self.cp_comm_type]
+                )
+                assert all(
+                    cp_comm_type is not None
+                    and cp_comm_type.replace("_", "").lower() == "allgather"
+                    for cp_comm_type in cp_comm_types
+                ), (
+                    "DSAttention context parallelism currently supports "
+                    "cp_comm_type=allgather only."
+                )
 
         if self.inference_fuse_tp_communication:
             assert self.transformer_impl == "inference_optimized", (
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index d02151066d9..28dc6301a53 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -2039,6 +2039,7 @@ def _add_network_size_args(parser):
         "output_layer_init_method",
         "embedding_init_method",
         "activation_func",
+        "experimental_attention_variant_loss_scale_func",
         # types affect docstring
         "pipeline_model_parallel_layout",
         "window_size",
diff --git a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py
index 9059d0157aa..810f48092ee 100644
--- a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py
+++ b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py
@@ -36,11 +36,52 @@ def dtype_tols(dtype):
 
 
 class FakeCPGroup:
+    def __init__(self, size=1, rank=0):
+        self._size = size
+        self._rank = rank
+
     def size(self):
-        return 1
+        return self._size
 
     def rank(self):
-        return 0
+        return self._rank
+
+
+class TestApplyRotaryPosEmbTHD:
+    def test_packed_freqs_returns_offset_mapped_output_for_context_parallel(self):
+        cp_group = FakeCPGroup(size=2, rank=0)
+        cu_seqlens = torch.tensor([0, 4, 8], dtype=torch.int32)
+        t = torch.randn(4, 2, 8)
+        freqs = torch.randn(8, 1, 1, 8)
+
+        out = rope_utils_module._apply_rotary_pos_emb_thd(t, cu_seqlens, freqs, cp_group=cp_group)
+
+        expected_freqs = torch.cat([freqs[0:1], freqs[3:4], freqs[4:5], freqs[7:8]], dim=0)
+        expected = rope_utils_module._apply_rotary_pos_emb_bshd(
+            t.unsqueeze(1), expected_freqs
+        ).squeeze(1)
+
+        torch.testing.assert_close(out, expected)
+
+    def test_max_seqlen_freqs_returns_sequence_mapped_output_for_context_parallel(self):
+        cp_group = FakeCPGroup(size=2, rank=1)
+        cu_seqlens = torch.tensor([0, 4, 8], dtype=torch.int32)
+        t = torch.randn(4, 2, 8)
+        freqs = torch.randn(4, 1, 1, 8)
+
+        out = rope_utils_module._apply_rotary_pos_emb_thd(t, cu_seqlens, freqs, cp_group=cp_group)
+
+        expected_freqs = torch.cat([freqs[1:2], freqs[2:3]], dim=0)
+        expected_slices = []
+        for x in torch.split(t, [2, 2]):
+            expected_slices.append(
+                rope_utils_module._apply_rotary_pos_emb_bshd(
+                    x.unsqueeze(1), expected_freqs
+                ).squeeze(1)
+            )
+        expected = torch.cat(expected_slices, dim=0)
+
+        torch.testing.assert_close(out, expected)
 
 
 def _test_fused_apply_mla_rope_for_q(input_format):
diff --git a/tests/unit_tests/models/test_experimental_attention_variant_module_specs.py b/tests/unit_tests/models/test_experimental_attention_variant_module_specs.py
index 7cc406a198a..0a454b5d7ff 100644
--- a/tests/unit_tests/models/test_experimental_attention_variant_module_specs.py
+++ b/tests/unit_tests/models/test_experimental_attention_variant_module_specs.py
@@ -319,12 +319,14 @@ def test_rejects_qk_l2_norm(self):
         with pytest.raises(AssertionError, match="qk_l2_norm is not supported"):
             get_dsa_module_spec_for_backend(cfg, backend=_make_backend())
 
-    def test_returns_mla_self_attention_spec(self):
-        """Verify the returned attention module is MLA self-attention with causal mask."""
-        from megatron.core.transformer.multi_latent_attention import MLASelfAttention
+    def test_returns_absorbed_mla_self_attention_spec(self):
+        """Verify the returned attention module is absorbed MLA with causal mask."""
+        from megatron.core.transformer.experimental_attention_variant.absorbed_mla import (
+            AbsorbedMLASelfAttention,
+        )
 
         spec = self._call()
-        assert spec.module is MLASelfAttention
+        assert spec.module is AbsorbedMLASelfAttention
         assert spec.params == {"attn_mask_type": AttnMaskType.causal}
         assert spec.metainfo == {"fuse_input_layernorm": False}
 
diff --git a/tests/unit_tests/pipeline_parallel/test_schedules.py b/tests/unit_tests/pipeline_parallel/test_schedules.py
index 7dbd9fb15b1..8444602481f 100644
--- a/tests/unit_tests/pipeline_parallel/test_schedules.py
+++ b/tests/unit_tests/pipeline_parallel/test_schedules.py
@@ -1,6 +1,7 @@
 # Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
 
 import os
+from types import SimpleNamespace
 
 import pytest
 import torch
@@ -78,6 +79,120 @@ def test_deallocate_output_tensor():
     assert out.nelement() == 6
 
 
+@pytest.mark.parametrize("calculate_per_token_loss,expected_scale", [(False, 6.0), (True, 3.0)])
+def test_dsa_indexer_loss_scale_matches_schedule_cp_scaling(
+    calculate_per_token_loss, expected_scale
+):
+    from megatron.core.transformer.experimental_attention_variant.dsa import (
+        DSAIndexerLossAutoScaler,
+    )
+
+    config = SimpleNamespace(
+        calculate_per_token_loss=calculate_per_token_loss,
+        experimental_attention_variant_loss_scale_func=DSAIndexerLossAutoScaler.set_loss_scale,
+        experimental_attention_variant='dsa',
+        grad_scale_func=lambda tensor: tensor * 3.0,
+        num_moe_experts=None,
+        mtp_num_layers=None,
+        timers=None,
+    )
+    forward_data_store = []
+
+    def loss_func(output_tensor):
+        return output_tensor.clone(), torch.tensor(4), {'loss_reduced': output_tensor.detach()}
+
+    DSAIndexerLossAutoScaler.main_loss_backward_scale = None
+    schedule.forward_step_calc_loss(
+        model=None,
+        output_tensor=torch.tensor(8.0),
+        loss_func=loss_func,
+        config=config,
+        vp_stage=None,
+        collect_non_loss_data=False,
+        num_microbatches=2,
+        forward_data_store=forward_data_store,
+        cp_group_size=4,
+        is_last_stage=True,
+    )
+
+    torch.testing.assert_close(
+        DSAIndexerLossAutoScaler.main_loss_backward_scale, torch.tensor([expected_scale])
+    )
+
+
+def test_dsa_indexer_loss_scale_accepts_dict_output_tensor():
+    from megatron.core.transformer.experimental_attention_variant.dsa import (
+        DSAIndexerLossAutoScaler,
+    )
+
+    config = SimpleNamespace(
+        calculate_per_token_loss=True,
+        experimental_attention_variant_loss_scale_func=DSAIndexerLossAutoScaler.set_loss_scale,
+        experimental_attention_variant='dsa',
+        grad_scale_func=lambda tensor: tensor * 5.0,
+        num_moe_experts=None,
+        mtp_num_layers=None,
+        timers=None,
+    )
+
+    forward_data_store = []
+
+    DSAIndexerLossAutoScaler.main_loss_backward_scale = None
+    schedule.forward_step_calc_loss(
+        model=None,
+        output_tensor={'loss': torch.tensor(8.0)},
+        loss_func=None,
+        config=config,
+        vp_stage=None,
+        collect_non_loss_data=False,
+        num_microbatches=2,
+        forward_data_store=forward_data_store,
+        cp_group_size=4,
+        is_last_stage=True,
+    )
+
+    assert len(forward_data_store) == 1
+    torch.testing.assert_close(forward_data_store[0]['loss'], torch.tensor(8.0))
+    torch.testing.assert_close(
+        DSAIndexerLossAutoScaler.main_loss_backward_scale, torch.tensor([5.0])
+    )
+
+
+def test_dsa_indexer_loss_scale_defaults_from_variant_without_mutating_config():
+    from megatron.core.transformer.experimental_attention_variant.dsa import (
+        DSAIndexerLossAutoScaler,
+    )
+
+    config = SimpleNamespace(
+        calculate_per_token_loss=True,
+        experimental_attention_variant_loss_scale_func=None,
+        experimental_attention_variant='dsa',
+        grad_scale_func=lambda tensor: tensor * 7.0,
+        num_moe_experts=None,
+        mtp_num_layers=None,
+        timers=None,
+    )
+
+    DSAIndexerLossAutoScaler.main_loss_backward_scale = None
+    schedule.forward_step_calc_loss(
+        model=None,
+        output_tensor=torch.tensor(8.0),
+        loss_func=None,
+        config=config,
+        vp_stage=None,
+        collect_non_loss_data=False,
+        num_microbatches=2,
+        forward_data_store=[],
+        cp_group_size=4,
+        is_last_stage=True,
+    )
+
+    assert config.experimental_attention_variant_loss_scale_func is None
+    torch.testing.assert_close(
+        DSAIndexerLossAutoScaler.main_loss_backward_scale, torch.tensor([7.0])
+    )
+
+
 @pytest.mark.internal
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
 @pytest.mark.parametrize(
diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py b/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py
index eb235501ad7..aa2c0a1479f 100644
--- a/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py
+++ b/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py
@@ -1,6 +1,7 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import random
+from types import SimpleNamespace
 from typing import List, Optional, Tuple
 
 import pytest
@@ -12,6 +13,9 @@
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
 from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.experimental_attention_variant import (
+    absorbed_mla as absorbed_mla_module,
+)
 from megatron.core.transformer.experimental_attention_variant.absorbed_mla import (
     AbsorbedMLASelfAttention,
     AbsorbedMLASelfAttentionSubmodules,
@@ -193,8 +197,7 @@ def get_absorbed_mla_submodules(
         linear_q_down_proj=linear_q_down_proj,
         linear_q_up_proj=backend.column_parallel_linear(),
         linear_kv_down_proj=linear_kv_down_proj,
-        linear_k_up_proj=backend.column_parallel_linear(),
-        linear_v_up_proj=backend.column_parallel_linear(),
+        linear_kv_up_proj=backend.column_parallel_linear(),
         core_attention=MockCoreAttention,
         linear_proj=backend.row_parallel_linear(),
         q_layernorm=qk_norm,
@@ -227,6 +230,96 @@ def get_mla_submodules(
     )
 
 
+def test_checkpointed_attention_forward_captures_metadata(monkeypatch):
+    """Optional metadata should stay in the closure instead of checkpoint tensor args."""
+
+    packed_seq_params = PackedSeqParams(qkv_format='thd')
+    checkpoint_args = None
+
+    def fake_checkpoint(run_function, distribute_saved_activations, *args):
+        nonlocal checkpoint_args
+        del distribute_saved_activations
+        checkpoint_args = args
+        assert all(torch.is_tensor(arg) for arg in args)
+        return run_function(*args)
+
+    class CoreAttention(torch.nn.Module):
+        def forward(self, query, key, value, attention_mask, **kwargs):
+            del query, key, value, attention_mask
+            assert kwargs["packed_seq_params"] is packed_seq_params
+            assert kwargs["position_ids"] is None
+            return kwargs["x"]
+
+    dummy_attention = type(
+        "DummyAttention",
+        (),
+        {"attn_mask_type": AttnMaskType.causal, "core_attention": CoreAttention()},
+    )()
+
+    monkeypatch.setattr(absorbed_mla_module.tensor_parallel, "checkpoint", fake_checkpoint)
+
+    hidden_states = torch.randn(4, 1, 8)
+    output = AbsorbedMLASelfAttention._checkpointed_attention_forward(
+        dummy_attention,
+        q_absorbed=torch.randn(4, 1, 2, 8),
+        k_compressed=torch.randn(4, 1, 1, 8),
+        hidden_states=hidden_states,
+        q_compressed=torch.randn(4, 1, 8),
+        attention_mask=torch.empty(1),
+        up_v_weight=torch.randn(2, 4, 4),
+        position_ids=None,
+        packed_seq_params=packed_seq_params,
+    )
+
+    assert checkpoint_args is not None
+    assert all(arg is not packed_seq_params for arg in checkpoint_args)
+    assert all(arg is not None for arg in checkpoint_args)
+    assert output is hidden_states
+
+
+def test_load_from_state_dict_combines_split_kv_up_projection(monkeypatch):
+    """Pre-refactor split K/V up-projection checkpoints should load into the combined layout."""
+
+    dummy_attention = object.__new__(AbsorbedMLASelfAttention)
+    dummy_attention.num_attention_heads_per_partition = 2
+    dummy_attention.config = SimpleNamespace(qk_head_dim=2, v_head_dim=3, kv_lora_rank=4)
+
+    prefix = "self_attention."
+    k_weight = torch.arange(2 * 2 * 4, dtype=torch.float32).view(2 * 2, 4)
+    v_weight = torch.arange(2 * 3 * 4, dtype=torch.float32).view(2 * 3, 4)
+    state_dict = {
+        f"{prefix}linear_k_up_proj.weight": k_weight.clone(),
+        f"{prefix}linear_v_up_proj.weight": v_weight.clone(),
+        f"{prefix}linear_k_up_proj._extra_state": torch.empty(0),
+        f"{prefix}linear_v_up_proj._extra_state": torch.empty(0),
+    }
+    captured_state_dict = {}
+
+    def fake_super_load(self, state_dict, *args, **kwargs):
+        del self, args, kwargs
+        captured_state_dict.update(state_dict)
+
+    monkeypatch.setattr(absorbed_mla_module.Attention, "_load_from_state_dict", fake_super_load)
+
+    AbsorbedMLASelfAttention._load_from_state_dict(
+        dummy_attention, state_dict, prefix, {}, True, [], [], []
+    )
+
+    expected_weight = (
+        torch.cat((k_weight.view(2, 2, 4), v_weight.view(2, 3, 4)), dim=1)
+        .contiguous()
+        .view(2 * (2 + 3), 4)
+    )
+    torch.testing.assert_close(
+        captured_state_dict[f"{prefix}linear_kv_up_proj.weight"], expected_weight
+    )
+    assert f"{prefix}linear_k_up_proj.weight" not in captured_state_dict
+    assert f"{prefix}linear_v_up_proj.weight" not in captured_state_dict
+    assert f"{prefix}linear_kv_up_proj._extra_state" in captured_state_dict
+    assert f"{prefix}linear_k_up_proj._extra_state" not in captured_state_dict
+    assert f"{prefix}linear_v_up_proj._extra_state" not in captured_state_dict
+
+
 @pytest.mark.parametrize("tp_cp", [[1, 1], [2, 1], [1, 2], [2, 2]])
 @pytest.mark.parametrize("qkv_format", ['sbhd', 'thd'])
 @pytest.mark.parametrize("down_proj_use_column_parallel", [False, True])
@@ -350,50 +443,17 @@ def _calculate_tensor_similarity(x, y):
     absorbed_grads = dict(absorbed_mla.named_parameters())
     standard_grads = dict(standard_mla.named_parameters())
 
-    # Map parameter names between absorbed and standard MLA
-    # Most parameters have the same name, except for K/V up proj
     for name, param in standard_grads.items():
-        if 'linear_kv_up_proj' in name:
-            # Special handling: combine k and v up proj grads from absorbed_mla
-            k_name = name.replace('linear_kv_up_proj', 'linear_k_up_proj')
-            v_name = name.replace('linear_kv_up_proj', 'linear_v_up_proj')
-
-            k_grad = absorbed_grads[k_name].grad
-            v_grad = absorbed_grads[v_name].grad
-
-            # Combine k and v grads (interleaved by head)
-            # k_grad: [n * qk_head_dim, kv_lora_rank]
-            # v_grad: [n * v_head_dim, kv_lora_rank]
-            # combined: [n * (qk_head_dim + v_head_dim), kv_lora_rank]
-            n_heads = absorbed_mla.num_attention_heads_per_partition
-            qk_head_dim = absorbed_mla.config.qk_head_dim
-            v_head_dim = absorbed_mla.config.v_head_dim
-            kv_lora_rank = absorbed_mla.config.kv_lora_rank
-
-            k_grad_3d = k_grad.view(n_heads, qk_head_dim, kv_lora_rank)
-            v_grad_3d = v_grad.view(n_heads, v_head_dim, kv_lora_rank)
-            combined_grad_3d = torch.cat([k_grad_3d, v_grad_3d], dim=1)
-            combined_grad = combined_grad_3d.view(-1, kv_lora_rank)
-
-            absorbed_grad_flat = combined_grad.flatten().float()
-            standard_grad_flat = param.grad.flatten().float()
-
-            cos_sim = torch.nn.functional.cosine_similarity(
-                absorbed_grad_flat.unsqueeze(0), standard_grad_flat.unsqueeze(0)
-            ).item()
-            assert cos_sim > 0.9999, f"name: {name}, cosine similarity = {cos_sim} < 0.9999"
-            assert _calculate_tensor_similarity(combined_grad, param.grad) > 0.9999
-        else:
-            absorbed_grad = absorbed_grads[name].grad
-            standard_grad = param.grad
+        absorbed_grad = absorbed_grads[name].grad
+        standard_grad = param.grad
 
-            absorbed_grad_flat = absorbed_grad.flatten().float()
-            standard_grad_flat = standard_grad.flatten().float()
+        absorbed_grad_flat = absorbed_grad.flatten().float()
+        standard_grad_flat = standard_grad.flatten().float()
 
-            cos_sim = torch.nn.functional.cosine_similarity(
-                absorbed_grad_flat.unsqueeze(0), standard_grad_flat.unsqueeze(0)
-            ).item()
-            assert cos_sim > 0.9999, f"name: {name}, cosine similarity = {cos_sim} < 0.9999"
-            assert _calculate_tensor_similarity(absorbed_grad, standard_grad) > 0.9999
+        cos_sim = torch.nn.functional.cosine_similarity(
+            absorbed_grad_flat.unsqueeze(0), standard_grad_flat.unsqueeze(0)
+        ).item()
+        assert cos_sim > 0.9999, f"name: {name}, cosine similarity = {cos_sim} < 0.9999"
+        assert _calculate_tensor_similarity(absorbed_grad, standard_grad) > 0.9999
 
     Utils.destroy_model_parallel()
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 757b9dd283a..82eb72ad71d 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
@@ -10,11 +10,16 @@
     get_dsa_module_spec_for_backend,
     get_experimental_attention_variant_module_spec,
 )
-from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
+from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
-from megatron.core.transformer import TransformerConfig
 from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.experimental_attention_variant import dsa_kernels
+from megatron.core.transformer.experimental_attention_variant.absorbed_mla import (
+    AbsorbedMLASelfAttention,
+    _apply_absorbed_v_up_projection,
+    _restore_packed_thd_batch_dim,
+)
 from megatron.core.transformer.experimental_attention_variant.dsa import (
     DSAIndexer,
     DSAIndexerLossAutoScaler,
@@ -22,22 +27,34 @@
     DSAttention,
     DSAttentionSubmodules,
     FusedDSAIndexerLoss,
-    _compute_index_scores,
+    _run_sparse_attention,
+    _validate_nonpacked_cp_uniform_length,
     compute_dsa_indexer_loss,
     fused_qk_topk_naive,
     rotate_activation,
+    unfused_dsa_fn,
+)
+from megatron.core.transformer.experimental_attention_variant.dsa_layout import (
+    build_packed_allgather_cp_query_positions_and_key_reorder,
+    build_zigzag_allgather_cp_key_reorder,
+    get_cp_positions_from_layout,
+)
+from megatron.core.transformer.experimental_attention_variant.dsa_masking import (
+    build_causal_mask_from_positions,
+    build_fused_indexer_varlen_bounds,
+    generate_varlen_mask_params,
+    scatter_topk_into_index_mask,
 )
-from megatron.core.transformer.multi_latent_attention import MLASelfAttention
 from megatron.core.transformer.transformer_config import MLATransformerConfig
 from tests.unit_tests.test_utilities import Utils
 
 try:
-    from fast_hadamard_transform import hadamard_transform as _hadamard_transform
+    from fast_hadamard_transform import hadamard_transform
 
     HAVE_HADAMARD = True
 except ImportError:
+    hadamard_transform = None
     HAVE_HADAMARD = False
-    _hadamard_transform = None
 
 
 def mock_hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
@@ -48,6 +65,126 @@ def mock_hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor
     return x * scale
 
 
+def _build_packed_causal_mask_for_test(
+    query_idx: torch.Tensor, key_idx: torch.Tensor, cu_seqlens: torch.Tensor
+) -> torch.Tensor:
+    """Build packed-sequence causal mask for tests."""
+    query_idx = query_idx.to(dtype=torch.int64)
+    key_idx = key_idx.to(dtype=torch.int64)
+    cu_seqlens = cu_seqlens.to(device=query_idx.device, dtype=torch.int64)
+
+    boundaries = cu_seqlens[1:]
+    query_seq_id = torch.searchsorted(boundaries, query_idx, right=True)
+    key_seq_id = torch.searchsorted(boundaries, key_idx, right=True)
+    valid = (query_seq_id.unsqueeze(-1) == key_seq_id.unsqueeze(0)) & (
+        key_idx.unsqueeze(0) <= query_idx.unsqueeze(-1)
+    )
+    mask = torch.zeros(
+        (query_idx.numel(), key_idx.numel()), dtype=torch.float32, device=query_idx.device
+    )
+    mask.masked_fill_(~valid, float("-inf"))
+    return mask
+
+
+def _assert_topk_indices_in_bounds_or_invalid(topk_indices: torch.Tensor, seqlen: int) -> None:
+    """Assert top-k indices are valid token ids or sanitized invalid slots."""
+    assert torch.all((topk_indices == -1) | ((topk_indices >= 0) & (topk_indices < seqlen)))
+
+
+def _assert_valid_topk_indices_unique(topk_indices: torch.Tensor) -> None:
+    """Assert non-negative top-k entries do not repeat within each row."""
+    sorted_indices = torch.sort(topk_indices, dim=-1).values
+    adjacent_valid = (sorted_indices[..., 1:] >= 0) & (sorted_indices[..., :-1] >= 0)
+    duplicate_valid = (sorted_indices[..., 1:] == sorted_indices[..., :-1]) & adjacent_valid
+    assert not torch.any(duplicate_valid)
+
+
+def _broadcast_from_global_rank0(tensor: torch.Tensor) -> torch.Tensor:
+    """Use one global test input across ranks before slicing it for TP comparisons."""
+    if torch.distributed.is_available() and torch.distributed.is_initialized():
+        torch.distributed.broadcast(tensor, src=0)
+    return tensor
+
+
+def _compute_sparse_topk_reference_loss(
+    *,
+    index_topk_scores: torch.Tensor,
+    topk_indices: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    softmax_scale: float,
+    loss_coeff: float,
+    query_valid_rows: torch.Tensor | None = None,
+    calculate_per_token_loss: bool = False,
+) -> torch.Tensor:
+    """Dense reference for sparse top-k indexer KL tests."""
+    sq, b, np, hn = query.size()
+    sk, bk, nk, hk = key.size()
+    assert bk == b and hk == hn
+    assert index_topk_scores.shape == topk_indices.shape
+    assert index_topk_scores.shape[:2] == (b, sq)
+    if nk != 1:
+        assert nk == np
+
+    idx_raw = topk_indices.to(dtype=torch.int64, device=query.device)
+    valid = idx_raw >= 0
+    idx = idx_raw.clamp(min=0)
+    topk = idx.size(-1)
+    target = torch.zeros((b, sq, topk), dtype=torch.float32, device=query.device)
+
+    for bi in range(b):
+        q_b = query[:, bi].permute(1, 0, 2).float()  # [np, sq, hn]
+        if nk == 1:
+            key_sel = key[:, bi, 0].float().index_select(0, idx[bi].reshape(-1))
+            key_sel = key_sel.view(sq, topk, hn)
+            logits = torch.einsum("hsd,skd->hsk", q_b, key_sel) * softmax_scale
+        else:
+            logits_per_head = []
+            for head in range(np):
+                key_sel = key[:, bi, head].float().index_select(0, idx[bi].reshape(-1))
+                key_sel = key_sel.view(sq, topk, hn)
+                logits_per_head.append((q_b[head].unsqueeze(1) * key_sel).sum(dim=-1))
+            logits = torch.stack(logits_per_head, dim=0) * softmax_scale
+
+        logits = logits.masked_fill(~valid[bi].unsqueeze(0), float("-inf"))
+        target[bi] = torch.softmax(logits, dim=-1, dtype=torch.float32).sum(dim=0)
+
+    target = target / target.sum(dim=-1, keepdim=True).clamp_min(1e-10)
+    index_logits = index_topk_scores.to(dtype=torch.float32, device=query.device)
+    index_logits = index_logits.masked_fill(~valid, float("-inf"))
+    no_valid_rows = ~valid.any(dim=-1, keepdim=True)
+    if no_valid_rows.any():
+        index_logits = index_logits.masked_fill(no_valid_rows.expand_as(index_logits), 0.0)
+    index_probs = torch.softmax(index_logits, dim=-1, dtype=torch.float32)
+    kl_per_row = (target * (torch.log(target + 1e-10) - torch.log(index_probs + 1e-10))).sum(dim=-1)
+
+    if query_valid_rows is not None:
+        query_valid_rows = query_valid_rows.to(device=query.device, dtype=torch.bool)
+        if query_valid_rows.ndim == 1:
+            query_valid_rows = query_valid_rows.view(1, sq).expand(b, sq)
+        kl_per_row = kl_per_row * query_valid_rows.to(dtype=kl_per_row.dtype)
+
+    if calculate_per_token_loss:
+        kl_div = kl_per_row.sum()
+    elif query_valid_rows is None:
+        kl_div = kl_per_row.mean()
+    else:
+        kl_div = kl_per_row.sum() / query_valid_rows.sum().to(dtype=torch.float32).clamp_min(1.0)
+    return kl_div * loss_coeff
+
+
+class _FakeCPGroup:
+    def __init__(self, size: int, rank: int = 0):
+        self._size = size
+        self._rank = rank
+
+    def size(self) -> int:
+        return self._size
+
+    def rank(self) -> int:
+        return self._rank
+
+
 @pytest.fixture(autouse=True)
 def patch_hadamard_if_needed():
     """Automatically patch hadamard_transform in dsa module if not installed."""
@@ -61,6 +198,691 @@ def patch_hadamard_if_needed():
         yield
 
 
+def test_dsa_kernel_backend_selects_optional_kernel_module():
+    """DSA kernel backend config should select one optional backend module."""
+
+    class Config:
+        attention_backend = "auto"
+        dsa_kernel_backend = "none"
+
+    config = Config()
+
+    assert dsa_kernels._get_backend_module_name(config) is None
+    assert not dsa_kernels.use_fused_dsa_kernels(config)
+
+    config.dsa_kernel_backend = "tilelang"
+    assert (
+        dsa_kernels._get_backend_module_name(config)
+        == "megatron.core.transformer.experimental_attention_variant.dsa_tilelang_kernels"
+    )
+    assert dsa_kernels.use_fused_dsa_kernels(config)
+
+    config.dsa_kernel_backend = "cudnn"
+    assert (
+        dsa_kernels._get_backend_module_name(config)
+        == "megatron.core.transformer.experimental_attention_variant.dsa_cudnn_kernels"
+    )
+
+    config.attention_backend = "unfused"
+    assert not dsa_kernels.use_fused_dsa_kernels(config)
+
+    config.attention_backend = "auto"
+    config.dsa_kernel_backend = "invalid"
+    with pytest.raises(ValueError, match="dsa_kernel_backend"):
+        dsa_kernels._get_backend_module_name(config)
+
+
+class TestDSACPPositionHelpers:
+    """Test helper utilities used for DSAttention context-parallel masking."""
+
+    def test_allgather_layout_positions(self):
+        """Allgather CP layout should map to zigzag query and global key positions."""
+        query_pos, key_pos = get_cp_positions_from_layout(
+            sq=4, skv=8, cp_size=2, cp_rank=1, cp_comm_type="allgather", device=torch.device("cpu")
+        )
+        assert query_pos.tolist() == [2, 3, 4, 5]
+        assert key_pos.tolist() == list(range(8))
+
+    def test_nonpacked_allgather_cp_layout_reorders_gathered_kv_to_global_order(self):
+        """Non-packed allgather-CP helper should mirror MCore zigzag local order."""
+        query_pos, _ = get_cp_positions_from_layout(
+            sq=4, skv=8, cp_size=2, cp_rank=0, cp_comm_type="allgather", device=torch.device("cpu")
+        )
+        key_reorder_idx = build_zigzag_allgather_cp_key_reorder(
+            sq=4, cp_size=2, device=torch.device("cpu")
+        )
+
+        assert query_pos.tolist() == [0, 1, 6, 7]
+
+        gathered_key_pos = torch.tensor([0, 1, 6, 7, 2, 3, 4, 5], dtype=torch.int64)
+        restored = gathered_key_pos.index_select(0, key_reorder_idx)
+        assert restored.tolist() == list(range(8))
+
+    def test_nonpacked_allgather_cp_rejects_uneven_rank_lengths(self, monkeypatch):
+        """Non-packed allgather CP requires uniform per-rank sequence lengths."""
+        local_lengths = [3, 5]
+        fake_cp_group = _FakeCPGroup(len(local_lengths))
+
+        monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
+        monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+
+        def _fake_all_gather(out, local_len, group=None):
+            del local_len, group
+            for i, tensor in enumerate(out):
+                tensor.copy_(
+                    torch.tensor([local_lengths[i]], dtype=tensor.dtype, device=tensor.device)
+                )
+
+        monkeypatch.setattr(torch.distributed, "all_gather", _fake_all_gather)
+
+        with pytest.raises(RuntimeError, match="uniform per-rank sequence lengths"):
+            _validate_nonpacked_cp_uniform_length(
+                sq=local_lengths[1],
+                skv=local_lengths[1],
+                cp_size=len(local_lengths),
+                cp_group=fake_cp_group,
+                device=torch.device("cpu"),
+            )
+
+    def test_position_based_causal_mask(self):
+        """Position-based causal mask should mask keys with strictly larger positions."""
+        query_pos = torch.tensor([0, 2], dtype=torch.int64)
+        key_pos = torch.tensor([0, 1, 2, 3], dtype=torch.int64)
+        mask = build_causal_mask_from_positions(query_pos, key_pos)
+        expected = torch.tensor(
+            [[0.0, float("-inf"), float("-inf"), float("-inf")], [0.0, 0.0, 0.0, float("-inf")]],
+            dtype=torch.float32,
+        )
+        torch.testing.assert_close(mask, expected, rtol=0, atol=0)
+
+    def test_packed_position_based_causal_mask(self):
+        """Packed causal mask should block cross-sequence attention using cu_seqlens boundaries."""
+        # Two packed sequences: [0,1,2] and [3,4]
+        cu_seqlens = torch.tensor([0, 3, 5], dtype=torch.int32)
+        query_idx = torch.tensor([1, 3, 4], dtype=torch.int64)
+        key_idx = torch.tensor([0, 1, 2, 3, 4], dtype=torch.int64)
+
+        mask = _build_packed_causal_mask_for_test(query_idx, key_idx, cu_seqlens)
+        expected = torch.tensor(
+            [
+                [0.0, 0.0, float("-inf"), float("-inf"), float("-inf")],
+                [float("-inf"), float("-inf"), float("-inf"), 0.0, float("-inf")],
+                [float("-inf"), float("-inf"), float("-inf"), 0.0, 0.0],
+            ],
+            dtype=torch.float32,
+        )
+        torch.testing.assert_close(mask, expected, rtol=0, atol=0)
+
+    def test_topk_uses_key_length(self):
+        """Top-k selection should be bounded by key length, not query length."""
+        sq, skv, bsz, nheads, dim = 4, 7, 1, 2, 8
+        topk = 6
+        q = torch.randn(sq, bsz, nheads, dim, dtype=torch.float32)
+        k = torch.randn(skv, bsz, dim, dtype=torch.float32)
+        weights = torch.randn(sq, bsz, nheads, dtype=torch.float32)
+
+        _, topk_indices = fused_qk_topk_naive(q, k, weights, topk, mask=None)
+        assert topk_indices.shape == (bsz, sq, topk)
+
+    def test_cp_packed_varlen_end_to_end_matches_dense_mask(self):
+        """CP+THD multi-sequence varlen path should match dense packed mask end-to-end."""
+        # Simulate cp_size=2 allgather layout with local query chunk and global keys.
+        cp_size, cp_rank = 2, 1
+        sq, skv = 4, 8
+        bsz, nheads, dim, vdim = 1, 2, 8, 6
+        topk = 4
+        softmax_scale = dim**-0.5
+
+        # Three packed sequences in global stream: [0,1,2], [3,4], [5,6,7]
+        cu_seqlens = torch.tensor([0, 3, 5, 8], dtype=torch.int32)
+        query_idx, key_idx = get_cp_positions_from_layout(
+            sq=sq,
+            skv=skv,
+            cp_size=cp_size,
+            cp_rank=cp_rank,
+            cp_comm_type="allgather",
+            device=torch.device("cpu"),
+        )
+
+        # Build varlen starts/ends for local query rows.
+        starts_all, ends_all = generate_varlen_mask_params(cu_seqlens.to(torch.int64))
+        starts = starts_all.index_select(0, query_idx)
+        ends = ends_all.index_select(0, query_idx)
+
+        q = torch.randn(sq, bsz, nheads, dim, dtype=torch.float32)
+        k_for_index = torch.randn(skv, bsz, dim, dtype=torch.float32)
+        weights = torch.randn(sq, bsz, nheads, dtype=torch.float32)
+        query = torch.randn(sq, bsz, nheads, dim, dtype=torch.float32)
+        key = torch.randn(skv, bsz, nheads, dim, dtype=torch.float32)
+        value = torch.randn(skv, bsz, nheads, vdim, dtype=torch.float32)
+
+        dense_mask = _build_packed_causal_mask_for_test(query_idx, key_idx, cu_seqlens)
+        _, dense_idx = fused_qk_topk_naive(q, k_for_index, weights, topk, mask=dense_mask)
+        out_dense = unfused_dsa_fn(query, key, value, dense_idx, softmax_scale, mask=dense_mask)
+
+        _, varlen_idx = fused_qk_topk_naive(
+            q,
+            k_for_index,
+            weights,
+            topk,
+            mask=None,
+            varlen_starts=starts,
+            varlen_ends=ends,
+            key_positions=key_idx,
+        )
+        out_varlen = unfused_dsa_fn(
+            query,
+            key,
+            value,
+            varlen_idx,
+            softmax_scale,
+            mask=None,
+            varlen_starts=starts,
+            varlen_ends=ends,
+            key_positions=key_idx,
+        )
+
+        torch.testing.assert_close(out_varlen, out_dense, rtol=0, atol=0)
+
+    def test_cp_packed_varlen_uneven_rank_lengths_matches_dense_mask(self, monkeypatch):
+        """CP+THD varlen path should match dense mask under uneven per-rank query lengths."""
+        # Simulate cp_size=2, cp_rank=1, local query lengths [3, 5].
+        cp_size, cp_rank = 2, 1
+        local_lengths = [3, 5]
+        sq, skv = local_lengths[cp_rank], sum(local_lengths)
+        bsz, nheads, dim, vdim = 1, 2, 8, 6
+        topk = 4
+        softmax_scale = dim**-0.5
+
+        fake_cp_group = _FakeCPGroup(cp_size)
+
+        monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
+        monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+
+        def _fake_all_gather(out, local_len, group=None):
+            del local_len, group
+            for i, tensor in enumerate(out):
+                tensor.copy_(
+                    torch.tensor([local_lengths[i]], dtype=tensor.dtype, device=tensor.device)
+                )
+
+        monkeypatch.setattr(torch.distributed, "all_gather", _fake_all_gather)
+
+        # Packed global stream has three sequences: [0,1], [2,3,4], [5,6,7]
+        cu_seqlens = torch.tensor([0, 2, 5, 8], dtype=torch.int32)
+        query_idx, key_idx = get_cp_positions_from_layout(
+            sq=sq,
+            skv=skv,
+            cp_size=cp_size,
+            cp_rank=cp_rank,
+            cp_comm_type="allgather",
+            device=torch.device("cpu"),
+            cp_group=fake_cp_group,
+        )
+        assert query_idx.tolist() == [3, 4, 5, 6, 7]
+
+        starts_all, ends_all = generate_varlen_mask_params(cu_seqlens.to(torch.int64))
+        starts = starts_all.index_select(0, query_idx)
+        ends = ends_all.index_select(0, query_idx)
+
+        q = torch.randn(sq, bsz, nheads, dim, dtype=torch.float32)
+        k_for_index = torch.randn(skv, bsz, dim, dtype=torch.float32)
+        weights = torch.randn(sq, bsz, nheads, dtype=torch.float32)
+        query = torch.randn(sq, bsz, nheads, dim, dtype=torch.float32)
+        key = torch.randn(skv, bsz, nheads, dim, dtype=torch.float32)
+        value = torch.randn(skv, bsz, nheads, vdim, dtype=torch.float32)
+
+        dense_mask = _build_packed_causal_mask_for_test(query_idx, key_idx, cu_seqlens)
+        _, dense_idx = fused_qk_topk_naive(q, k_for_index, weights, topk, mask=dense_mask)
+        out_dense = unfused_dsa_fn(query, key, value, dense_idx, softmax_scale, mask=dense_mask)
+
+        _, varlen_idx = fused_qk_topk_naive(
+            q,
+            k_for_index,
+            weights,
+            topk,
+            mask=None,
+            varlen_starts=starts,
+            varlen_ends=ends,
+            key_positions=key_idx,
+        )
+        out_varlen = unfused_dsa_fn(
+            query,
+            key,
+            value,
+            varlen_idx,
+            softmax_scale,
+            mask=None,
+            varlen_starts=starts,
+            varlen_ends=ends,
+            key_positions=key_idx,
+        )
+
+        torch.testing.assert_close(out_varlen, out_dense, rtol=0, atol=0)
+
+    def test_packed_allgather_cp_layout_reorders_gathered_kv_to_global_order(self):
+        """Packed allgather-CP helper should mirror zigzag local order and restore global KV order."""
+        cu_seqlens = torch.tensor([0, 8, 16], dtype=torch.int32)
+
+        query_pos, key_reorder_idx = build_packed_allgather_cp_query_positions_and_key_reorder(
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_kv=cu_seqlens,
+            cp_size=2,
+            cp_rank=0,
+            device=torch.device("cpu"),
+        )
+
+        assert query_pos.tolist() == [0, 1, 6, 7, 8, 9, 14, 15]
+
+        gathered_key_pos = torch.tensor(
+            [0, 1, 6, 7, 8, 9, 14, 15, 2, 3, 4, 5, 10, 11, 12, 13], dtype=torch.int64
+        )
+        restored = gathered_key_pos.index_select(0, key_reorder_idx)
+        assert restored.tolist() == list(range(16))
+
+    def test_cp_packed_zigzag_varlen_matches_dense_mask(self):
+        """Packed zigzag CP query positions + gathered-KV reorder should match dense masking."""
+        cp_size, cp_rank = 2, 1
+        cu_seqlens = torch.tensor([0, 8, 16], dtype=torch.int32)
+        bsz, nheads, dim, vdim = 1, 2, 8, 6
+        topk = 4
+        softmax_scale = dim**-0.5
+
+        query_pos, key_reorder_idx = build_packed_allgather_cp_query_positions_and_key_reorder(
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_kv=cu_seqlens,
+            cp_size=cp_size,
+            cp_rank=cp_rank,
+            device=torch.device("cpu"),
+        )
+        sq, skv = query_pos.numel(), int(cu_seqlens[-1].item())
+        key_pos = torch.arange(skv, dtype=torch.int64)
+
+        gathered_key_order = torch.empty_like(key_reorder_idx)
+        gathered_key_order[key_reorder_idx] = torch.arange(skv, dtype=torch.int64)
+
+        starts_all, ends_all = generate_varlen_mask_params(cu_seqlens.to(torch.int64))
+        starts = starts_all.index_select(0, query_pos)
+        ends = ends_all.index_select(0, query_pos)
+
+        q = torch.randn(sq, bsz, nheads, dim, dtype=torch.float32)
+        k_for_index_global = torch.randn(skv, bsz, dim, dtype=torch.float32)
+        weights = torch.randn(sq, bsz, nheads, dtype=torch.float32)
+        query = torch.randn(sq, bsz, nheads, dim, dtype=torch.float32)
+        key_global = torch.randn(skv, bsz, nheads, dim, dtype=torch.float32)
+        value_global = torch.randn(skv, bsz, nheads, vdim, dtype=torch.float32)
+
+        dense_mask = _build_packed_causal_mask_for_test(query_pos, key_pos, cu_seqlens)
+        _, dense_idx = fused_qk_topk_naive(q, k_for_index_global, weights, topk, mask=dense_mask)
+        out_dense = unfused_dsa_fn(
+            query, key_global, value_global, dense_idx, softmax_scale, mask=dense_mask
+        )
+
+        k_for_index_gathered = k_for_index_global.index_select(0, gathered_key_order)
+        key_gathered = key_global.index_select(0, gathered_key_order)
+        value_gathered = value_global.index_select(0, gathered_key_order)
+
+        k_for_index_reordered = k_for_index_gathered.index_select(0, key_reorder_idx)
+        key_reordered = key_gathered.index_select(0, key_reorder_idx)
+        value_reordered = value_gathered.index_select(0, key_reorder_idx)
+
+        _, varlen_idx = fused_qk_topk_naive(
+            q,
+            k_for_index_reordered,
+            weights,
+            topk,
+            mask=None,
+            varlen_starts=starts,
+            varlen_ends=ends,
+            key_positions=key_pos,
+        )
+        out_varlen = unfused_dsa_fn(
+            query,
+            key_reordered,
+            value_reordered,
+            varlen_idx,
+            softmax_scale,
+            mask=None,
+            varlen_starts=starts,
+            varlen_ends=ends,
+            key_positions=key_pos,
+        )
+
+        torch.testing.assert_close(out_varlen, out_dense, rtol=0, atol=0)
+
+    def test_unfused_dsa_allows_delayed_backward_after_same_shape_reuse(self):
+        """Unfused DSA should not mutate tensors saved by earlier forward graphs."""
+        torch.manual_seed(123)
+        sq, bsz, nheads, dim, vdim = 4, 1, 2, 3, 2
+        topk_indices = (
+            torch.arange(sq, dtype=torch.int64).view(1, 1, sq).expand(bsz, sq, sq).contiguous()
+        )
+
+        query = torch.randn(sq, bsz, nheads, dim, dtype=torch.float32, requires_grad=True)
+        key = torch.randn(sq, bsz, nheads, dim, dtype=torch.float32, requires_grad=True)
+        value = torch.randn(sq, bsz, nheads, vdim, dtype=torch.float32, requires_grad=True)
+
+        out1 = unfused_dsa_fn(query, key, value, topk_indices, dim**-0.5)
+        out2 = unfused_dsa_fn(query, key, value, topk_indices, dim**-0.5)
+        (out1.square().sum() + out2.square().sum()).backward()
+
+        assert query.grad is not None and torch.isfinite(query.grad).all()
+        assert key.grad is not None and torch.isfinite(key.grad).all()
+        assert value.grad is not None and torch.isfinite(value.grad).all()
+
+    def test_unfused_dsa_all_invalid_topk_rows_keep_gradients_finite(self):
+        """Rows with no valid sparse entries should avoid NaNs in autograd."""
+        torch.manual_seed(123)
+        sq, bsz, nheads, dim, vdim = 4, 1, 2, 3, 2
+        topk_indices = torch.full((bsz, sq, 3), -1, dtype=torch.int64)
+
+        query = torch.randn(sq, bsz, nheads, dim, dtype=torch.float32, requires_grad=True)
+        key = torch.randn(sq, bsz, nheads, dim, dtype=torch.float32, requires_grad=True)
+        value = torch.randn(sq, bsz, nheads, vdim, dtype=torch.float32, requires_grad=True)
+
+        out = unfused_dsa_fn(query, key, value, topk_indices, dim**-0.5)
+        out.square().sum().backward()
+
+        assert torch.isfinite(out).all()
+        assert query.grad is not None and torch.isfinite(query.grad).all()
+        assert key.grad is not None and torch.isfinite(key.grad).all()
+        assert value.grad is not None and torch.isfinite(value.grad).all()
+
+    def test_fused_bounds_disable_on_per_batch_mask_mismatch(self):
+        """Fused bounds should disable when batched masks are not identical."""
+        sq, skv, bsz = 5, 7, 2
+        base_mask = torch.triu(
+            torch.full((sq, skv), float("-inf"), dtype=torch.float32), diagonal=1
+        )
+        mask = base_mask.unsqueeze(0).expand(bsz, -1, -1).clone()
+        out = build_fused_indexer_varlen_bounds(
+            sq=sq,
+            skv=skv,
+            device=mask.device,
+            mask=mask,
+            varlen_starts=None,
+            varlen_ends=None,
+            key_positions=None,
+        )
+        assert out is not None
+
+        # Change one batch mask so masks are no longer identical.
+        mask[1, 0, 0] = float("-inf")
+        out_mismatch = build_fused_indexer_varlen_bounds(
+            sq=sq,
+            skv=skv,
+            device=mask.device,
+            mask=mask,
+            varlen_starts=None,
+            varlen_ends=None,
+            key_positions=None,
+        )
+        assert out_mismatch is None
+
+    def test_scatter_topk_chunked_matches_manual_with_negative_indices(self):
+        """Chunked top-k scatter should match manual behavior for -1 invalid indices."""
+        b, sq, skv = 2, 4, 6
+        topk_indices = torch.tensor(
+            [
+                [[0, 2, -1], [1, -1, -1], [2, 4, 5], [3, -1, 0]],
+                [[5, 4, 1], [0, -1, 2], [3, -1, -1], [1, 2, 3]],
+            ],
+            dtype=torch.int32,
+        )
+        got = torch.full((b, sq, skv), float("-inf"), dtype=torch.float32)
+        scatter_topk_into_index_mask(got, topk_indices, seq_chunk_size=2)
+
+        expected = torch.full((b, sq, skv), float("-inf"), dtype=torch.float32)
+        topk_i64 = topk_indices.to(torch.int64)
+        valid = topk_i64 >= 0
+        b_idx, q_idx, t_idx = torch.where(valid)
+        k_idx = topk_i64[b_idx, q_idx, t_idx]
+        expected[b_idx, q_idx, k_idx] = 0.0
+
+        assert torch.equal(got, expected)
+
+
+class TestDSAAbsorbedParityCPU:
+    """CPU parity tests for absorbed DSA rewrite."""
+
+    def test_absorbed_path_matches_non_absorbed_output(self):
+        """Absorbed attention + up_v projection should match non-absorbed attention output."""
+        torch.manual_seed(1234)
+
+        sq, skv, bsz, nheads = 6, 6, 1, 3
+        qk_dim, qk_pos_dim = 5, 2
+        kv_lora_rank, vdim = 4, 3
+        softmax_scale = (qk_dim + qk_pos_dim) ** -0.5
+
+        # Build synthetic tensors consistent with the absorbed rewrite equations.
+        q_no_pe = torch.randn(sq, bsz, nheads, qk_dim, dtype=torch.float32)
+        q_pos = torch.randn(sq, bsz, nheads, qk_pos_dim, dtype=torch.float32)
+        kv_latent = torch.randn(skv, bsz, kv_lora_rank, dtype=torch.float32)
+        k_pos_shared = torch.randn(skv, bsz, 1, qk_pos_dim, dtype=torch.float32)
+
+        up_k_weight = torch.randn(nheads, qk_dim, kv_lora_rank, dtype=torch.float32)
+        up_v_weight = torch.randn(nheads, vdim, kv_lora_rank, dtype=torch.float32)
+
+        # Non-absorbed tensors.
+        query_non_abs = torch.cat([q_no_pe, q_pos], dim=-1).contiguous()
+        k_no_pe = torch.einsum("sbk,hqk->sbhq", kv_latent, up_k_weight)
+        key_non_abs = torch.cat([k_no_pe, k_pos_shared.expand(-1, -1, nheads, -1)], dim=-1)
+        value_non_abs = torch.einsum("sbk,hvk->sbhv", kv_latent, up_v_weight).contiguous()
+
+        # Absorbed tensors.
+        q_content_abs = torch.einsum("sbhq,hqk->sbhk", q_no_pe, up_k_weight)
+        query_abs = torch.cat([q_content_abs, q_pos], dim=-1).contiguous()
+        key_abs = torch.cat([kv_latent.unsqueeze(2), k_pos_shared], dim=-1).contiguous()
+
+        # Use full-key support and causal masking in both paths.
+        topk_indices = (
+            torch.arange(skv, dtype=torch.int64).view(1, 1, skv).expand(bsz, sq, skv).contiguous()
+        )
+        causal_mask = torch.triu(
+            torch.full((sq, skv), float("-inf"), dtype=torch.float32), diagonal=1
+        )
+
+        out_non_abs = unfused_dsa_fn(
+            query_non_abs, key_non_abs, value_non_abs, topk_indices, softmax_scale, mask=causal_mask
+        )
+        config = type(
+            "Config", (), {"kv_lora_rank": kv_lora_rank, "attention_backend": "unfused"}
+        )()
+        out_abs = _run_sparse_attention(
+            absorbed_mla=True,
+            query=query_abs,
+            key=key_abs,
+            value=None,
+            up_v_weight=up_v_weight,
+            topk_indices=topk_indices,
+            softmax_scale=softmax_scale,
+            config=config,
+            mask=causal_mask,
+            varlen_starts=None,
+            varlen_ends=None,
+            key_positions=None,
+        )
+
+        torch.testing.assert_close(out_abs, out_non_abs, rtol=1e-4, atol=1e-5)
+
+    def test_absorbed_path_requires_up_v_weight(self):
+        """Absorbed attention must project latent output back to value head dim."""
+        sq, bsz, nheads = 2, 1, 2
+        kv_lora_rank, qk_pos_dim = 4, 2
+        config = type(
+            "Config", (), {"kv_lora_rank": kv_lora_rank, "attention_backend": "unfused"}
+        )()
+
+        query = torch.randn(sq, bsz, nheads, kv_lora_rank + qk_pos_dim)
+        key = torch.randn(sq, bsz, 1, kv_lora_rank + qk_pos_dim)
+        topk_indices = torch.arange(sq, dtype=torch.int64).view(1, 1, sq).expand(bsz, sq, sq)
+
+        with pytest.raises(RuntimeError, match="requires up_v_weight"):
+            _run_sparse_attention(
+                absorbed_mla=True,
+                query=query,
+                key=key,
+                value=None,
+                up_v_weight=None,
+                topk_indices=topk_indices,
+                softmax_scale=1.0,
+                config=config,
+                mask=None,
+                varlen_starts=None,
+                varlen_ends=None,
+                key_positions=None,
+            )
+
+
+class TestAbsorbedMLAPackedTHDShape:
+    """CPU tests for packed-THD absorbed MLA output rank handling."""
+
+    def test_restore_packed_thd_batch_dim_when_core_output_is_2d(self):
+        hidden_states = torch.empty(7, 1, 16)
+        core_attn_out = torch.empty(7, 16)
+        packed_seq_params = PackedSeqParams(qkv_format='thd')
+
+        restored = _restore_packed_thd_batch_dim(core_attn_out, hidden_states, packed_seq_params)
+
+        assert restored.shape == (7, 1, 16)
+
+    def test_restore_packed_thd_batch_dim_keeps_already_normalized_output(self):
+        hidden_states = torch.empty(7, 1, 16)
+        core_attn_out = torch.empty(7, 1, 16)
+        packed_seq_params = PackedSeqParams(qkv_format='thd')
+
+        restored = _restore_packed_thd_batch_dim(core_attn_out, hidden_states, packed_seq_params)
+
+        assert restored is core_attn_out
+        assert restored.shape == hidden_states.shape
+
+
+class TestAbsorbedMLAVUpProjection:
+    """CPU tests for absorbed MLA V-up projection decisions."""
+
+    def test_projection_applies_when_core_did_not_consume_weight_even_if_sizes_match(self):
+        torch.manual_seed(123)
+        num_heads, kv_lora_rank, v_head_dim = 2, 3, 3
+        core_attn_out = torch.randn(5, 1, num_heads * kv_lora_rank)
+        v_up_weight = torch.randn(num_heads, v_head_dim, kv_lora_rank)
+
+        projected = _apply_absorbed_v_up_projection(
+            core_attn_out,
+            v_up_weight,
+            num_attention_heads_per_partition=num_heads,
+            kv_lora_rank=kv_lora_rank,
+            v_head_dim=v_head_dim,
+            core_consumed_v_up_projection=False,
+        )
+        expected = core_attn_out.view(5, 1, num_heads, kv_lora_rank)
+        expected = torch.einsum("...nc,ndc->...nd", expected, v_up_weight)
+        expected = expected.contiguous().view(5, 1, -1)
+
+        torch.testing.assert_close(projected, expected, rtol=0, atol=0)
+
+    def test_projection_skips_when_core_consumed_weight_even_if_sizes_match(self):
+        num_heads, kv_lora_rank, v_head_dim = 2, 3, 3
+        core_attn_out = torch.randn(5, 1, num_heads * v_head_dim)
+        v_up_weight = torch.randn(num_heads, v_head_dim, kv_lora_rank)
+
+        projected = _apply_absorbed_v_up_projection(
+            core_attn_out,
+            v_up_weight,
+            num_attention_heads_per_partition=num_heads,
+            kv_lora_rank=kv_lora_rank,
+            v_head_dim=v_head_dim,
+            core_consumed_v_up_projection=True,
+        )
+
+        assert projected is core_attn_out
+
+
+class TestDSAIndexerLossRowMaskCPU:
+    """CPU tests for packed-row masking in DSA indexer loss."""
+
+    @staticmethod
+    def _fake_pg_collection():
+        class _FakeTP:
+            @staticmethod
+            def size():
+                return 1
+
+        class _FakeCollection:
+            tp = _FakeTP()
+
+        return _FakeCollection()
+
+    def test_dense_indexer_loss_ignores_padded_rows(self):
+        index_scores = torch.tensor([[[2.0, float("-inf")], [0.1, 0.9]]], dtype=torch.float32)
+        topk_indices = torch.tensor([[[0, 1], [1, 0]]], dtype=torch.int64)
+        query = torch.tensor([[[[1.0, 0.0]]], [[[0.0, 1.0]]]], dtype=torch.float32)
+        key = torch.tensor([[[[1.0, 0.0]]], [[[0.0, 1.0]]]], dtype=torch.float32)
+        mask = torch.tensor([[0.0, float("-inf")], [0.0, 0.0]], dtype=torch.float32)
+
+        masked_loss = compute_dsa_indexer_loss(
+            index_scores=index_scores.clone(),
+            topk_indices=topk_indices,
+            query=query,
+            key=key,
+            softmax_scale=1.0,
+            loss_coeff=1.0,
+            sparse_loss=False,
+            pg_collection=self._fake_pg_collection(),
+            mask=mask,
+            query_valid_rows=torch.tensor([True, False], dtype=torch.bool),
+        )
+        trimmed_loss = compute_dsa_indexer_loss(
+            index_scores=index_scores[:, :1, :].clone(),
+            topk_indices=topk_indices[:, :1, :].clone(),
+            query=query[:1].clone(),
+            key=key,
+            softmax_scale=1.0,
+            loss_coeff=1.0,
+            sparse_loss=False,
+            pg_collection=self._fake_pg_collection(),
+            mask=mask[:1],
+        )
+
+        torch.testing.assert_close(masked_loss, trimmed_loss)
+
+    def test_sparse_indexer_loss_ignores_padded_rows(self):
+        index_topk_scores = torch.tensor([[[2.0, float("-inf")], [0.9, 0.1]]], dtype=torch.float32)
+        topk_indices = torch.tensor([[[0, 1], [1, 0]]], dtype=torch.int64)
+        query = torch.tensor([[[[1.0, 0.0]]], [[[0.0, 1.0]]]], dtype=torch.float32)
+        key = torch.tensor([[[[1.0, 0.0]]], [[[0.0, 1.0]]]], dtype=torch.float32)
+
+        masked_loss = _compute_sparse_topk_reference_loss(
+            index_topk_scores=index_topk_scores.clone(),
+            topk_indices=topk_indices,
+            query=query,
+            key=key,
+            softmax_scale=1.0,
+            loss_coeff=1.0,
+            query_valid_rows=torch.tensor([True, False], dtype=torch.bool),
+        )
+        trimmed_loss = _compute_sparse_topk_reference_loss(
+            index_topk_scores=index_topk_scores[:, :1, :].clone(),
+            topk_indices=topk_indices[:, :1, :].clone(),
+            query=query[:1].clone(),
+            key=key,
+            softmax_scale=1.0,
+            loss_coeff=1.0,
+        )
+
+        torch.testing.assert_close(masked_loss, trimmed_loss)
+
+    def test_naive_topk_masks_all_invalid_slots_with_minus_one(self):
+        q = torch.tensor([[[[1.0]]]], dtype=torch.float32)
+        k = torch.tensor([[[1.0]], [[0.0]], [[0.0]]], dtype=torch.float32)
+        weights = torch.tensor([[[1.0]]], dtype=torch.float32)
+        mask = torch.tensor([[0.0, float("-inf"), float("-inf")]], dtype=torch.float32)
+
+        _, topk_indices = fused_qk_topk_naive(q=q, k=k, weights=weights, index_topk=3, mask=mask)
+
+        expected = torch.tensor([[[0, -1, -1]]], dtype=torch.int64)
+        torch.testing.assert_close(topk_indices, expected)
+
+
 class TestRotateActivation:
     """Test rotate_activation function."""
 
@@ -212,6 +1034,76 @@ def test_dsa_indexer_loss_sparse(self, seqlen_and_topk):
         assert loss_sparse >= 0
         assert loss_dense >= 0
 
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    def test_sparse_varlen_empty_rows_are_finite(self, seqlen_and_topk):
+        """Sparse varlen rows with no valid keys should not produce NaN gradients."""
+        del seqlen_and_topk
+        seqlen = 3
+        batch_size = 1
+        num_heads = 2
+        head_dim = 4
+        index_n_heads = 2
+        index_head_dim = 4
+
+        q = torch.randn(
+            seqlen,
+            batch_size,
+            index_n_heads,
+            index_head_dim,
+            dtype=torch.float32,
+            device="cuda",
+            requires_grad=True,
+        )
+        weights = torch.randn(
+            seqlen,
+            batch_size,
+            index_n_heads,
+            dtype=torch.float32,
+            device="cuda",
+            requires_grad=True,
+        )
+        k = torch.randn(
+            seqlen,
+            batch_size,
+            index_head_dim,
+            dtype=torch.float32,
+            device="cuda",
+            requires_grad=True,
+        )
+        query = torch.randn(seqlen, batch_size, num_heads, head_dim, dtype=torch.bfloat16).cuda()
+        key = torch.randn(seqlen, batch_size, num_heads, head_dim, dtype=torch.bfloat16).cuda()
+
+        varlen_starts = torch.tensor([0, 0, 2], dtype=torch.int64, device="cuda")
+        varlen_ends = torch.tensor([1, 0, 3], dtype=torch.int64, device="cuda")
+        key_positions = torch.arange(seqlen, dtype=torch.int64, device="cuda")
+        query_valid_rows = torch.tensor([[True, False, True]], dtype=torch.bool, device="cuda")
+
+        _, loss = FusedDSAIndexerLoss.apply(
+            q,
+            weights,
+            k,
+            query,
+            key,
+            1.0,
+            2,
+            0.01,
+            None,
+            True,
+            self.pg_collection,
+            varlen_starts,
+            varlen_ends,
+            key_positions,
+            query_valid_rows,
+            False,
+            False,
+        )
+
+        assert torch.isfinite(loss)
+        loss.backward()
+        assert torch.isfinite(q.grad).all()
+        assert torch.isfinite(weights.grad).all()
+        assert torch.isfinite(k.grad).all()
+
 
 class TestDSAIndexerLossAutoScaler:
     """Test DSAIndexerLossAutoScaler autograd function."""
@@ -248,8 +1140,9 @@ def test_backward_pass(self):
         dummy_input.requires_grad_(True)
         indexer_loss = dummy_input.mean()
 
-        # Set loss scale
-        scale = torch.tensor(2.0).cuda()
+        # Set loss scale. The schedule can supply this from CPU while the
+        # indexer loss graph is on CUDA.
+        scale = torch.tensor(2.0)
         DSAIndexerLossAutoScaler.set_loss_scale(scale)
 
         # Apply the autograd function
@@ -274,6 +1167,13 @@ def test_backward_pass(self):
             atol=0,
         ), f"Gradient should be scaled by loss scale, expected {expected_grad_per_element}, got {dummy_input.grad[0].item()}"
 
+    def test_set_loss_scale_requires_tensor(self):
+        """set_loss_scale has the same tensor-only contract as other auxiliary loss scalers."""
+        DSAIndexerLossAutoScaler.main_loss_backward_scale = torch.tensor(1.0)
+        with pytest.raises(TypeError, match="requires a torch.Tensor"):
+            DSAIndexerLossAutoScaler.set_loss_scale(1.0)
+        DSAIndexerLossAutoScaler.main_loss_backward_scale = None
+
 
 class TestFusedDSAIndexerLossGradient:
     """Test that FusedDSAIndexerLoss manual backward matches autograd backward."""
@@ -337,10 +1237,9 @@ def test_fused_indexer_loss_gradient_matches_autograd(self):
                 )
 
                 # Method 1: Autograd (reference)
-                index_scores_ref = _compute_index_scores(q_ref, weights_ref, k_ref)
-                index_scores_masked = index_scores_ref + mask.unsqueeze(0)
-                topk_k = min(index_topk, seqlen)
-                topk_indices = index_scores_masked.topk(topk_k, dim=-1)[1]
+                index_scores_masked, topk_indices = fused_qk_topk_naive(
+                    q_ref, k_ref, weights_ref, index_topk, mask=mask
+                )
 
                 loss_ref = compute_dsa_indexer_loss(
                     index_scores=index_scores_masked,
@@ -375,6 +1274,9 @@ def test_fused_indexer_loss_gradient_matches_autograd(self):
                     mask,
                     sparse_loss,
                     self.pg_collection,
+                    None,
+                    None,
+                    None,
                 )
                 loss_fused.backward()
 
@@ -472,6 +1374,9 @@ def test_fused_indexer_loss_gradient_tp_consistency(self):
                 mask,
                 sparse_loss,
                 pg_collection_tp1,
+                None,
+                None,
+                None,
             )
             loss_tp1.backward()
 
@@ -528,6 +1433,9 @@ def test_fused_indexer_loss_gradient_tp_consistency(self):
                     mask,
                     sparse_loss,
                     pg_collection_tpn,
+                    None,
+                    None,
+                    None,
                 )
                 loss_tpn.backward()
 
@@ -597,6 +1505,7 @@ def setup_method(self, request):
             use_cpu_initialization=True,
             bf16=True,
             params_dtype=torch.bfloat16,
+            layernorm_epsilon=1e-5,
             # MLA specific configs
             q_lora_rank=64,
             kv_lora_rank=64,
@@ -610,6 +1519,7 @@ def setup_method(self, request):
             dsa_indexer_n_heads=8,
             dsa_indexer_head_dim=64,
             dsa_indexer_topk=cls.index_topk,
+            dsa_indexer_k_norm_epsilon=1e-6,
         )
 
         # Create indexer submodules spec
@@ -636,6 +1546,57 @@ def test_dsa_indexer_constructor(self, seqlen):
         assert self.indexer.index_n_heads == 8
         assert self.indexer.index_head_dim == 64
         assert self.indexer.index_topk == 32
+        assert self.indexer.k_norm.eps == pytest.approx(1e-6)
+
+    @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."""
+        del seqlen
+        captured = {}
+
+        def _fake_apply_rotary_pos_emb(x, rotary_pos_emb, **kwargs):
+            captured["mla_rotary_interleaved"] = kwargs["mla_rotary_interleaved"]
+            return x
+
+        self.indexer.config.dsa_indexer_rope_interleaved = interleaved
+
+        x = torch.randn(
+            2, 1, self.indexer.index_n_heads, self.indexer.index_head_dim, dtype=torch.bfloat16
+        )
+        rotary_pos_emb = torch.randn(2, 1, 1, self.config.qk_pos_emb_head_dim, dtype=torch.bfloat16)
+
+        with patch(
+            "megatron.core.transformer.experimental_attention_variant.dsa.apply_rotary_pos_emb",
+            side_effect=_fake_apply_rotary_pos_emb,
+        ):
+            out = self.indexer._apply_rope(x, rotary_pos_emb, mscale=1.0)
+
+        assert captured["mla_rotary_interleaved"] is interleaved
+        assert out.shape == x.shape
+
+    @pytest.mark.parametrize("rotate_activation_enabled", [False, True])
+    def test_dsa_indexer_rotate_activation_follows_config(self, seqlen, rotate_activation_enabled):
+        """Ensure indexer Hadamard rotation can be disabled for GLM5-compatible scoring."""
+        del seqlen
+        self.indexer.config.dsa_indexer_rotate_activation = rotate_activation_enabled
+
+        self.indexer.cuda()
+        x = torch.randn(2, 1, self.config.hidden_size, dtype=torch.bfloat16).cuda()
+        qr = torch.randn(2, 1, self.config.q_lora_rank, dtype=torch.bfloat16).cuda()
+
+        with (
+            patch.object(self.indexer, "_apply_rope", side_effect=lambda t, *args, **kwargs: t),
+            patch(
+                "megatron.core.transformer.experimental_attention_variant.dsa.rotate_activation",
+                side_effect=lambda t: t,
+            ) as rotate_mock,
+        ):
+            q, k, _ = self.indexer.forward_before_topk(x, qr)
+
+        expected_calls = 2 if rotate_activation_enabled else 0
+        assert rotate_mock.call_count == expected_calls
+        assert q.shape[-1] == self.indexer.index_head_dim
+        assert k.shape[-1] == self.indexer.index_head_dim
 
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     def test_dsa_indexer_forward(self, seqlen):
@@ -654,12 +1615,9 @@ def test_dsa_indexer_forward(self, seqlen):
         # Check output shape
         assert topk_indices.shape == (batch_size, seqlen, min(self.config.dsa_indexer_topk, seqlen))
         assert topk_indices.dtype == torch.long
-        assert torch.all((topk_indices >= 0) & (topk_indices < seqlen))
+        _assert_topk_indices_in_bounds_or_invalid(topk_indices, seqlen)
         # Make sure no duplicate indices are selected
-        assert torch.all(
-            torch.sort(topk_indices, dim=-1).values[:, :, 1:]
-            != torch.sort(topk_indices, dim=-1).values[:, :, :-1]
-        )
+        _assert_valid_topk_indices_unique(topk_indices)
 
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     def test_dsa_indexer_forward_with_scores(self, seqlen):
@@ -680,13 +1638,40 @@ def test_dsa_indexer_forward_with_scores(self, seqlen):
         assert topk_indices.shape == (batch_size, seqlen, min(self.config.dsa_indexer_topk, seqlen))
         assert index_scores.dtype == torch.float32
         assert topk_indices.dtype == torch.long
-        assert torch.all((topk_indices >= 0) & (topk_indices < seqlen))
+        _assert_topk_indices_in_bounds_or_invalid(topk_indices, seqlen)
         # Make sure no duplicate indices are selected
-        assert torch.all(
-            torch.sort(topk_indices, dim=-1).values[:, :, 1:]
-            != torch.sort(topk_indices, dim=-1).values[:, :, :-1]
+        _assert_valid_topk_indices_unique(topk_indices)
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    def test_dsa_indexer_forward_with_scores_packed_thd(self, seqlen):
+        """Test indexer forward_with_scores works with packed THD inputs."""
+        batch_size = 1
+        self.indexer.cuda()
+
+        x = torch.randn(seqlen, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda()
+        qr = torch.randn(seqlen, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda()
+
+        cu_seqlens = torch.tensor([0, seqlen], dtype=torch.int32, device=x.device)
+        packed_seq_params = PackedSeqParams(
+            qkv_format="thd",
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_kv=cu_seqlens,
+            max_seqlen_q=seqlen,
+            max_seqlen_kv=seqlen,
+        )
+        token_idx = torch.arange(seqlen, dtype=torch.int64, device=x.device)
+        mask = _build_packed_causal_mask_for_test(token_idx, token_idx, cu_seqlens)
+
+        index_scores, topk_indices = self.indexer.forward_with_scores(
+            x, qr, mask=mask, packed_seq_params=packed_seq_params
         )
 
+        assert index_scores.shape == (batch_size, seqlen, seqlen)
+        assert topk_indices.shape == (batch_size, seqlen, min(self.config.dsa_indexer_topk, seqlen))
+        assert index_scores.dtype == torch.float32
+        assert topk_indices.dtype == torch.long
+        _assert_topk_indices_in_bounds_or_invalid(topk_indices, seqlen)
+
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     def test_dsa_indexer_with_mask(self, seqlen):
         """Test indexer with attention mask."""
@@ -782,6 +1767,177 @@ def test_dsa_constructor(self):
         assert isinstance(self.sparse_attention, DSAttention)
         assert hasattr(self.sparse_attention, 'indexer')
         assert isinstance(self.sparse_attention.indexer, DSAIndexer)
+        assert self.config.experimental_attention_variant_loss_scale_func is None
+
+    def test_unfused_backend_skips_full_fused_attention(self, monkeypatch):
+        """attention_backend=unfused must bypass optional full fused DSA kernels."""
+        seq_len = 4
+        batch_size = 1
+        num_heads = self.config.num_attention_heads
+        head_dim = self.config.hidden_size // num_heads
+
+        def _unexpected_fused_attention(**_kwargs):
+            raise AssertionError(
+                "full fused DSA backend should not run for attention_backend=unfused"
+            )
+
+        def _fake_forward_before_topk(_x, _qr, _packed_seq_params):
+            q_indexer = torch.randn(seq_len, batch_size, 2, 4)
+            k_indexer = torch.randn(seq_len, batch_size, 4)
+            weights = torch.ones(seq_len, batch_size, 2)
+            return q_indexer, k_indexer, weights
+
+        expected_output = torch.randn(seq_len, batch_size, self.config.hidden_size)
+
+        def _fake_run_sparse_attention(**_kwargs):
+            return expected_output
+
+        monkeypatch.setattr(self.config, "attention_backend", "unfused")
+        monkeypatch.setattr(
+            "megatron.core.transformer.experimental_attention_variant.dsa."
+            "dsa_kernels.run_fused_dsa_attention",
+            _unexpected_fused_attention,
+        )
+        monkeypatch.setattr(
+            self.sparse_attention.indexer, "forward_before_topk", _fake_forward_before_topk
+        )
+        monkeypatch.setattr(
+            "megatron.core.transformer.experimental_attention_variant.dsa._run_sparse_attention",
+            _fake_run_sparse_attention,
+        )
+
+        was_training = self.sparse_attention.training
+        self.sparse_attention.eval()
+        try:
+            output = self.sparse_attention(
+                query=torch.randn(seq_len, batch_size, num_heads, head_dim),
+                key=torch.randn(seq_len, batch_size, num_heads, head_dim),
+                value=torch.randn(seq_len, batch_size, num_heads, head_dim),
+                x=torch.randn(seq_len, batch_size, self.config.hidden_size),
+                qr=torch.randn(seq_len, batch_size, self.config.q_lora_rank),
+                attention_mask=None,
+                attn_mask_type=AttnMaskType.causal,
+            )
+        finally:
+            self.sparse_attention.train(was_training)
+
+        assert output is expected_output
+
+    def test_disabled_indexer_loss_can_use_full_fused_attention(self, monkeypatch):
+        """Full fused DSA attention forward can run when indexer loss is disabled."""
+        seq_len = 4
+        batch_size = 1
+        num_heads = self.config.num_attention_heads
+        head_dim = self.config.hidden_size // num_heads
+
+        def _fake_forward_before_topk(_x, _qr, _packed_seq_params):
+            q_indexer = torch.randn(seq_len, batch_size, 2, 4)
+            k_indexer = torch.randn(seq_len, batch_size, 4)
+            weights = torch.ones(seq_len, batch_size, 2)
+            return q_indexer, k_indexer, weights
+
+        expected_output = torch.randn(seq_len, batch_size, self.config.hidden_size)
+        seen = {}
+
+        def _fake_fused_attention(**kwargs):
+            seen["loss_coeff"] = kwargs["loss_coeff"]
+            return expected_output, torch.zeros((), dtype=torch.float32)
+
+        monkeypatch.setattr(self.config, "attention_backend", "auto")
+        monkeypatch.setattr(self.config, "dsa_kernel_backend", "cudnn")
+        monkeypatch.setattr(self.config, "dsa_indexer_loss_coeff", 0.0)
+        monkeypatch.setattr(
+            "megatron.core.transformer.experimental_attention_variant.dsa."
+            "dsa_kernels.run_fused_dsa_attention",
+            _fake_fused_attention,
+        )
+        monkeypatch.setattr(
+            self.sparse_attention.indexer, "forward_before_topk", _fake_forward_before_topk
+        )
+
+        was_training = self.sparse_attention.training
+        self.sparse_attention.train()
+        try:
+            output = self.sparse_attention(
+                query=torch.randn(seq_len, batch_size, num_heads, head_dim),
+                key=torch.randn(seq_len, batch_size, num_heads, head_dim),
+                value=torch.randn(seq_len, batch_size, num_heads, head_dim),
+                x=torch.randn(seq_len, batch_size, self.config.hidden_size),
+                qr=torch.randn(seq_len, batch_size, self.config.q_lora_rank),
+                attention_mask=None,
+                attn_mask_type=AttnMaskType.causal,
+            )
+        finally:
+            self.sparse_attention.train(was_training)
+
+        assert output is expected_output
+        assert seen["loss_coeff"] == 0.0
+
+    def test_packed_dense_indexer_loss_uses_local_varlen_on_fused_path(self, monkeypatch):
+        """Packed dense indexer loss should keep local varlen and be owned by the backend."""
+        seq_len = 4
+        key_seq_len = seq_len * 2
+        batch_size = 1
+        num_heads = self.config.num_attention_heads
+        head_dim = self.config.hidden_size // num_heads
+        seen = {}
+
+        def _fake_forward_before_topk(_x, _qr, _packed_seq_params):
+            q_indexer = torch.randn(seq_len, batch_size, 2, 4)
+            k_indexer = torch.randn(key_seq_len, batch_size, 4)
+            weights = torch.ones(seq_len, batch_size, 2)
+            return q_indexer, k_indexer, weights
+
+        expected_output = torch.randn(seq_len, batch_size, self.config.hidden_size)
+
+        def _fake_run_fused_attention(**kwargs):
+            seen["fused_loss_coeff"] = kwargs["loss_coeff"]
+            seen["fused_sparse_loss"] = kwargs["sparse_loss"]
+            seen["use_local_indexer_varlen"] = kwargs["use_local_indexer_varlen"]
+            return expected_output, torch.zeros((), dtype=torch.float32)
+
+        monkeypatch.setattr(self.config, "attention_backend", "auto")
+        monkeypatch.setattr(self.config, "dsa_kernel_backend", "cudnn")
+        monkeypatch.setattr(self.config, "dsa_indexer_use_sparse_loss", False)
+        monkeypatch.setattr(self.sparse_attention, "cp_comm_type", "allgather")
+        monkeypatch.setattr(self.sparse_attention.indexer.pg_collection, "cp", _FakeCPGroup(2))
+        monkeypatch.setattr(
+            self.sparse_attention.indexer, "forward_before_topk", _fake_forward_before_topk
+        )
+        monkeypatch.setattr(
+            "megatron.core.transformer.experimental_attention_variant.dsa."
+            "dsa_kernels.run_fused_dsa_attention",
+            _fake_run_fused_attention,
+        )
+
+        was_training = self.sparse_attention.training
+        self.sparse_attention.train()
+        cu_seqlens = torch.tensor([0, key_seq_len], dtype=torch.int32)
+        packed_seq_params = PackedSeqParams(
+            qkv_format="thd",
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_kv=cu_seqlens,
+            max_seqlen_q=key_seq_len,
+            max_seqlen_kv=key_seq_len,
+        )
+        try:
+            output = self.sparse_attention(
+                query=torch.randn(seq_len, batch_size, num_heads, head_dim),
+                key=torch.randn(key_seq_len, batch_size, num_heads, head_dim),
+                value=torch.randn(key_seq_len, batch_size, num_heads, head_dim),
+                x=torch.randn(seq_len, batch_size, self.config.hidden_size),
+                qr=torch.randn(seq_len, batch_size, self.config.q_lora_rank),
+                attention_mask=None,
+                attn_mask_type=AttnMaskType.causal,
+                packed_seq_params=packed_seq_params,
+            )
+        finally:
+            self.sparse_attention.train(was_training)
+
+        torch.testing.assert_close(output, expected_output)
+        assert seen["fused_loss_coeff"] == self.config.dsa_indexer_loss_coeff
+        assert seen["fused_sparse_loss"] is False
+        assert seen["use_local_indexer_varlen"] is True
 
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     def test_dsa_forward(self):
@@ -934,8 +2090,7 @@ def test_dsa_topk_selection(self):
             )
 
         # Check that topk_indices are valid
-        assert torch.all(topk_indices >= 0)
-        assert torch.all(topk_indices < seq_len)
+        _assert_topk_indices_in_bounds_or_invalid(topk_indices, seq_len)
         assert topk_indices.shape[2] == min(self.config.dsa_indexer_topk, seq_len)
 
 
@@ -1309,21 +2464,15 @@ def test_dsa_forward_consistency(self):
             num_heads = config_tp1.num_attention_heads
             head_dim = config_tp1.hidden_size // num_heads
 
-            query_input = (
-                torch.randn(seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16)
-                .cuda()
-                .requires_grad_(True)
-            )
-            key_input = (
-                torch.randn(seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16)
-                .cuda()
-                .requires_grad_(True)
-            )
-            value_input = (
-                torch.randn(seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16)
-                .cuda()
-                .requires_grad_(True)
-            )
+            query_input = torch.randn(
+                seq_len, batch_size, num_heads, head_dim, dtype=torch.float32
+            ).cuda()
+            key_input = torch.randn(
+                seq_len, batch_size, num_heads, head_dim, dtype=torch.float32
+            ).cuda()
+            value_input = torch.randn(
+                seq_len, batch_size, num_heads, head_dim, dtype=torch.float32
+            ).cuda()
             x_input = torch.randn(
                 seq_len, batch_size, config_tp1.hidden_size, dtype=torch.bfloat16
             ).cuda()
@@ -1332,6 +2481,15 @@ def test_dsa_forward_consistency(self):
             ).cuda()
             attention_mask = torch.ones(batch_size, 1, seq_len, seq_len, dtype=torch.bool).cuda()
             attention_mask = torch.tril(attention_mask)
+            query_input = _broadcast_from_global_rank0(query_input)
+            key_input = _broadcast_from_global_rank0(key_input)
+            value_input = _broadcast_from_global_rank0(value_input)
+            x_input = _broadcast_from_global_rank0(x_input)
+            qr_input = _broadcast_from_global_rank0(qr_input)
+            attention_mask = _broadcast_from_global_rank0(attention_mask)
+            query_input.requires_grad_(True)
+            key_input.requires_grad_(True)
+            value_input.requires_grad_(True)
 
             sparse_attention_tp1.train()
             output_tp1 = sparse_attention_tp1(
@@ -1359,6 +2517,16 @@ def test_dsa_forward_consistency(self):
                 value_input.grad.clone().cpu(),
                 num_heads,
                 head_dim,
+                query_input.detach().clone(),
+                key_input.detach().clone(),
+                value_input.detach().clone(),
+                x_input.detach().clone(),
+                qr_input.detach().clone(),
+                attention_mask.detach().clone(),
+                {
+                    name: tensor.detach().clone()
+                    for name, tensor in sparse_attention_tp1.indexer.state_dict().items()
+                },
             )
 
         Utils.destroy_model_parallel()
@@ -1383,6 +2551,13 @@ def test_dsa_forward_consistency(self):
                         value_tp1_grad,
                         num_heads,
                         head_dim,
+                        query_input_base,
+                        key_input_base,
+                        value_input_base,
+                        x_input_base,
+                        qr_input_base,
+                        attention_mask_base,
+                        indexer_tp1_state,
                     ) = baselines[use_sparse_indexer_loss]
 
                     config_tpn = self._create_config(
@@ -1395,27 +2570,15 @@ def test_dsa_forward_consistency(self):
                     sparse_attention_tpn = self._create_sparse_attention(
                         config_tpn, pg_collection_tpn
                     ).cuda()
+                    sparse_attention_tpn.indexer.load_state_dict(indexer_tp1_state)
                     tag = f"[TP={tensor_model_parallel_size}, SP={sequence_parallel}, sparse={use_sparse_indexer_loss}]"
 
-                    query_input_tpn = torch.randn(
-                        seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16
-                    ).cuda()
-                    key_input_tpn = torch.randn(
-                        seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16
-                    ).cuda()
-                    value_input_tpn = torch.randn(
-                        seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16
-                    ).cuda()
-                    x_input_tpn = torch.randn(
-                        seq_len, batch_size, config_tpn.hidden_size, dtype=torch.bfloat16
-                    ).cuda()
-                    qr_input_tpn = torch.randn(
-                        seq_len, batch_size, config_tpn.q_lora_rank, dtype=torch.bfloat16
-                    ).cuda()
-                    attention_mask_tpn = torch.ones(
-                        batch_size, 1, seq_len, seq_len, dtype=torch.bool
-                    ).cuda()
-                    attention_mask_tpn = torch.tril(attention_mask_tpn)
+                    query_input_tpn = query_input_base.detach().clone()
+                    key_input_tpn = key_input_base.detach().clone()
+                    value_input_tpn = value_input_base.detach().clone()
+                    x_input_tpn = x_input_base.detach().clone()
+                    qr_input_tpn = qr_input_base.detach().clone()
+                    attention_mask_tpn = attention_mask_base.detach().clone()
 
                     tp_rank = parallel_state.get_tensor_model_parallel_rank()
                     if sequence_parallel:
@@ -1459,9 +2622,13 @@ def test_dsa_forward_consistency(self):
                         output_tpn, group=pg_collection_tpn.tp
                     )
                     assert output_tpn_gathered.shape == output_tp1.shape
-                    assert torch.allclose(
-                        output_tpn_gathered.detach(), output_tp1, rtol=0, atol=0
-                    ), f"{tag} Sparse attention outputs mismatch vs TP=1"
+                    torch.testing.assert_close(
+                        output_tpn_gathered.detach(),
+                        output_tp1,
+                        rtol=1e-5,
+                        atol=1e-5,
+                        msg=f"{tag} Sparse attention outputs mismatch vs TP=1",
+                    )
 
                     for name, param in sparse_attention_tpn.indexer.named_parameters():
                         if param.grad is not None and name in indexer_tp1_grads:
@@ -1480,15 +2647,27 @@ def test_dsa_forward_consistency(self):
                         value_tpn.grad.reshape(sq, b, nh * hd), group=pg_collection_tpn.tp
                     ).reshape(sq, b, num_heads, hd)
 
-                    assert torch.allclose(
-                        query_grad_gathered.cpu(), query_tp1_grad, rtol=0, atol=0
-                    ), f"{tag} Query gradient mismatch vs TP=1"
-                    assert torch.allclose(
-                        key_grad_gathered.cpu(), key_tp1_grad, rtol=0, atol=0
-                    ), f"{tag} Key gradient mismatch vs TP=1"
-                    assert torch.allclose(
-                        value_grad_gathered.cpu(), value_tp1_grad, rtol=0, atol=0
-                    ), f"{tag} Value gradient mismatch vs TP=1"
+                    torch.testing.assert_close(
+                        query_grad_gathered.cpu(),
+                        query_tp1_grad,
+                        rtol=1e-5,
+                        atol=1e-5,
+                        msg=f"{tag} Query gradient mismatch vs TP=1",
+                    )
+                    torch.testing.assert_close(
+                        key_grad_gathered.cpu(),
+                        key_tp1_grad,
+                        rtol=1e-5,
+                        atol=1e-5,
+                        msg=f"{tag} Key gradient mismatch vs TP=1",
+                    )
+                    torch.testing.assert_close(
+                        value_grad_gathered.cpu(),
+                        value_tp1_grad,
+                        rtol=1e-5,
+                        atol=1e-5,
+                        msg=f"{tag} Value gradient mismatch vs TP=1",
+                    )
 
             Utils.destroy_model_parallel()
 
@@ -1645,9 +2824,16 @@ def test_get_experimental_attention_variant_module_spec_dsa(self):
         """get_experimental_attention_variant_module_spec dispatches to DSA for variant='dsa'."""
         config = self._make_dsa_config(experimental_attention_variant="dsa")
         spec = get_experimental_attention_variant_module_spec(config)
-        assert spec.module == MLASelfAttention
+        assert spec.module == AbsorbedMLASelfAttention
         assert spec.submodules.core_attention.module == DSAttention
 
+    def test_dsa_cp_requires_allgather_cp_comm_type(self):
+        """DSA context parallelism should fail early for unsupported CP communication."""
+        with pytest.raises(AssertionError, match="allgather"):
+            self._make_dsa_config(
+                experimental_attention_variant="dsa", context_parallel_size=2, cp_comm_type="p2p"
+            )
+
     def test_get_dsa_module_spec_for_backend(self):
         """get_dsa_module_spec_for_backend returns the correct full spec structure."""
         from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider
@@ -1655,7 +2841,7 @@ def test_get_dsa_module_spec_for_backend(self):
         config = self._make_dsa_config()
         backend = TESpecProvider()
         spec = get_dsa_module_spec_for_backend(config, backend=backend)
-        assert spec.module == MLASelfAttention
+        assert spec.module == AbsorbedMLASelfAttention
         assert spec.submodules.core_attention.module == DSAttention
         assert spec.submodules.core_attention.submodules.indexer.module == DSAIndexer
         assert spec.params["attn_mask_type"] == AttnMaskType.causal