[2/2] Add TileLang fused DSA kernels support with THD and CP & Clean up

megatron/core/transformer/experimental_attention_variant/dsa_cudnn_kernels.py

@@ -725,42 +725,15 @@ def _indexer_topk_multi_packed_cp_thd(
         raise RuntimeError("packed CP cuDNN THD indexer requires positive maximum sequence lengths")
 
     segment_divisor = 2 * cp_size
-    if sk % segment_divisor != 0:
-        raise RuntimeError(f"packed CP key length must be divisible by {segment_divisor}, got {sk}")
-
     device = q_bshd.device
-    cu_q = packed_cu_seqlens_q.to(device=device, dtype=torch.int64).contiguous()
-    cu_k = packed_cu_seqlens_k.to(device=device, dtype=torch.int64).contiguous()
-    q_lengths = cu_q[1:] - cu_q[:-1]
-    k_lengths = cu_k[1:] - cu_k[:-1]
-    q_half = q_lengths // segment_divisor
-    k_half = k_lengths // segment_divisor
-    segment_q_lengths = torch.stack((q_half, q_half), dim=1).reshape(-1)
-    segment_k_lengths = torch.stack(
-        ((cp_rank + 1) * k_half, k_lengths - cp_rank * k_half), dim=1
-    ).reshape(-1)
-
-    zero_i32 = torch.zeros(1, dtype=torch.int32, device=device)
-    segment_cu_q = torch.cat(
-        (zero_i32, segment_q_lengths.cumsum(dim=0, dtype=torch.int32))
-    ).contiguous()
-    segment_cu_k = torch.cat(
-        (zero_i32, segment_k_lengths.cumsum(dim=0, dtype=torch.int32))
-    ).contiguous()
-
-    segment_key_starts = cu_k[:-1].repeat_interleave(2)
-    total_segment_k = sk + sk // segment_divisor
-    segment_ids = torch.repeat_interleave(
-        torch.arange(segment_k_lengths.numel(), device=device),
-        segment_k_lengths,
-        output_size=total_segment_k,
-    )
-    segment_offsets = torch.arange(total_segment_k, device=device, dtype=torch.int64)
-    segment_offsets -= torch.repeat_interleave(
-        segment_cu_k[:-1].to(dtype=torch.int64), segment_k_lengths, output_size=total_segment_k
+    layout = dsa_layout.build_packed_cp_indexer_layout(
+        packed_cu_seqlens_q.to(device=device),
+        packed_cu_seqlens_k.to(device=device),
+        cp_size=cp_size,
+        cp_rank=cp_rank,
+        key_size=sk,
     )
-    source_indices = segment_key_starts.index_select(0, segment_ids) + segment_offsets
-    segmented_k = k_bshd[0].index_select(0, source_indices).contiguous()
+    segmented_k = k_bshd[0].index_select(0, layout.source_indices).contiguous()
 
     max_segment_q = packed_max_seqlen_q // segment_divisor
     max_k_half = packed_max_seqlen_k // segment_divisor
@@ -771,8 +744,8 @@ def _indexer_topk_multi_packed_cp_thd(
         w_bsh[0],
         ratio=_INDEXER_RATIO,
         sm_scale=_INDEXER_SOFTMAX_SCALE,
-        cu_seqlens_q=segment_cu_q,
-        cu_seqlens_k=segment_cu_k,
+        cu_seqlens_q=layout.segment_cu_q.to(dtype=torch.int32),
+        cu_seqlens_k=layout.segment_cu_k.to(dtype=torch.int32),
         max_seqlen_q=max_segment_q,
         max_seqlen_k=max_segment_k,
     )["scores"]
@@ -1043,11 +1016,8 @@ def _sort_valid_topk_indices_by_index(topk_indices: Tensor, topk_length: Tensor,
     """Canonicalize consumed top-K indices while keeping ignored suffix slots invalid."""
     positions = _trailing_positions(topk_indices)
     valid = positions < topk_length.unsqueeze(-1)
-    sort_key = torch.where(valid, topk_indices, torch.full_like(topk_indices, sk))
-    order = sort_key.argsort(dim=-1)
-    sorted_indices = torch.gather(topk_indices, dim=-1, index=order)
-    sorted_valid = torch.gather(valid.expand_as(topk_indices), dim=-1, index=order)
-    return sorted_indices.masked_fill(~sorted_valid, -1).contiguous()
+    sorted_indices, _ = dsa_masking.sort_topk_by_index(topk_indices, valid, sk=sk)
+    return sorted_indices
 
 
 def _sort_valid_topk_indices_and_scores_by_index(
@@ -1056,14 +1026,15 @@ def _sort_valid_topk_indices_and_scores_by_index(
     """Sort valid top-K indices and keep the selected score payload aligned."""
     positions = _trailing_positions(topk_indices)
     valid = positions < topk_length.unsqueeze(-1)
-    sort_key = torch.where(valid, topk_indices, torch.full_like(topk_indices, sk))
-    order = sort_key.argsort(dim=-1)
-    sorted_indices = torch.gather(topk_indices, dim=-1, index=order)
-    sorted_scores = torch.gather(topk_scores, dim=-1, index=order)
-    sorted_valid = torch.gather(valid.expand_as(topk_indices), dim=-1, index=order)
-    sorted_indices = sorted_indices.masked_fill(~sorted_valid, -1)
-    sorted_scores = sorted_scores.masked_fill(~sorted_valid, torch.finfo(torch.float32).min)
-    return sorted_indices.contiguous(), sorted_scores.contiguous()
+    sorted_indices, sorted_scores = dsa_masking.sort_topk_by_index(
+        topk_indices,
+        valid,
+        sk=sk,
+        topk_scores=topk_scores,
+        invalid_score=torch.finfo(torch.float32).min,
+    )
+    assert sorted_scores is not None
+    return sorted_indices, sorted_scores
 
 
 def _prepare_attention_topk_indices(topk_indices: Tensor, sk: int) -> Tuple[Tensor, Tensor]:

megatron/core/transformer/experimental_attention_variant/dsa_layout.py

@@ -2,6 +2,7 @@
 
 """Layout helpers for DeepSeek sparse attention."""
 
+from dataclasses import dataclass
 from typing import Optional, Tuple
 
 import torch
@@ -10,6 +11,8 @@
 from megatron.core.utils import get_pg_size
 
 __all__ = [
+    "PackedCPIndexerLayout",
+    "build_packed_cp_indexer_layout",
     "build_packed_allgather_cp_local_positions",
     "build_packed_allgather_cp_query_positions_and_key_reorder",
     "build_zigzag_allgather_cp_key_reorder",
@@ -22,6 +25,93 @@
 ]
 
 
+@dataclass(frozen=True)
+class PackedCPIndexerLayout:
+    """Segment metadata shared by packed-CP DSA indexer backends."""
+
+    segment_q_lengths: torch.Tensor
+    segment_k_lengths: torch.Tensor
+    segment_cu_q: torch.Tensor
+    segment_cu_k: torch.Tensor
+    segment_key_starts: torch.Tensor
+    source_indices: torch.Tensor
+
+
+def build_packed_cp_indexer_layout(
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_kv: torch.Tensor,
+    *,
+    cp_size: int,
+    cp_rank: int,
+    key_size: int,
+    local_key_layout: bool = False,
+) -> PackedCPIndexerLayout:
+    """Build packed-CP front/back segment metadata for fused DSA indexers.
+
+    ``local_key_layout`` describes the single-sequence optimization where the
+    key tensor contains only this CP rank's local front/back chunks. Otherwise,
+    ``key_size`` is the globally ordered packed key length.
+    """
+    if cp_size <= 1 or not 0 <= cp_rank < cp_size:
+        raise RuntimeError("packed CP indexer layout requires a valid CP rank and cp_size > 1")
+    if cu_seqlens_q.shape != cu_seqlens_kv.shape or cu_seqlens_q.numel() < 2:
+        raise RuntimeError("packed CP indexer layout requires matching non-empty q/k cu_seqlens")
+
+    device = cu_seqlens_q.device
+    cu_q = cu_seqlens_q.to(device=device, dtype=torch.int64).contiguous()
+    cu_k = cu_seqlens_kv.to(device=device, dtype=torch.int64).contiguous()
+    segment_divisor = 2 * cp_size
+
+    if local_key_layout:
+        if cu_q.numel() != 2 or key_size % 2 != 0:
+            raise RuntimeError(
+                "local-key packed CP indexer layout requires one sequence and even key rows"
+            )
+        half = key_size // 2
+        segment_q_lengths = torch.full((2,), half, dtype=torch.int64, device=device)
+        segment_k_lengths = torch.tensor((half, key_size), dtype=torch.int64, device=device)
+        segment_key_starts = torch.zeros(2, dtype=torch.int64, device=device)
+        total_segment_k = key_size + half
+    else:
+        if key_size % segment_divisor != 0:
+            raise RuntimeError(
+                f"packed CP key length must be divisible by {segment_divisor}, got {key_size}"
+            )
+        q_lengths = cu_q[1:] - cu_q[:-1]
+        k_lengths = cu_k[1:] - cu_k[:-1]
+        q_half = q_lengths // segment_divisor
+        k_half = k_lengths // segment_divisor
+        segment_q_lengths = torch.stack((q_half, q_half), dim=1).reshape(-1)
+        segment_k_lengths = torch.stack(
+            ((cp_rank + 1) * k_half, k_lengths - cp_rank * k_half), dim=1
+        ).reshape(-1)
+        segment_key_starts = cu_k[:-1].repeat_interleave(2)
+        total_segment_k = key_size + key_size // segment_divisor
+
+    zero = torch.zeros(1, dtype=torch.int64, device=device)
+    segment_cu_q = torch.cat((zero, segment_q_lengths.cumsum(dim=0))).contiguous()
+    segment_cu_k = torch.cat((zero, segment_k_lengths.cumsum(dim=0))).contiguous()
+
+    segment_ids = torch.repeat_interleave(
+        torch.arange(segment_k_lengths.numel(), device=device),
+        segment_k_lengths,
+        output_size=total_segment_k,
+    )
+    segment_offsets = torch.arange(total_segment_k, device=device, dtype=torch.int64)
+    segment_offsets -= torch.repeat_interleave(
+        segment_cu_k[:-1], segment_k_lengths, output_size=total_segment_k
+    )
+    source_indices = segment_key_starts.index_select(0, segment_ids) + segment_offsets
+    return PackedCPIndexerLayout(
+        segment_q_lengths=segment_q_lengths,
+        segment_k_lengths=segment_k_lengths,
+        segment_cu_q=segment_cu_q,
+        segment_cu_k=segment_cu_k,
+        segment_key_starts=segment_key_starts,
+        source_indices=source_indices,
+    )
+
+
 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:

megatron/core/transformer/experimental_attention_variant/dsa_masking.py

@@ -29,6 +29,7 @@
     "prepare_additive_mask",
     "prepare_sparse_mask_context",
     "scatter_topk_into_index_mask",
+    "sort_topk_by_index",
 ]
 
 
@@ -98,6 +99,37 @@ def build_valid_mask_from_starts_ends(
     )
 
 
+def sort_topk_by_index(
+    topk_indices: torch.Tensor,
+    valid_mask: torch.Tensor,
+    *,
+    sk: int,
+    topk_scores: Optional[torch.Tensor] = None,
+    invalid_score: float = float("-inf"),
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """Sort valid top-k slots by key index while preserving aligned scores.
+
+    Backends define validity explicitly: TileLang uses ``index >= 0`` sentinels,
+    while cuDNN consumes a compact prefix described by ``topk_length``.
+    """
+    if valid_mask.dtype != torch.bool or valid_mask.shape != topk_indices.shape:
+        raise ValueError("valid_mask must be boolean and match topk_indices")
+    if topk_scores is not None and topk_scores.shape != topk_indices.shape:
+        raise ValueError("topk_scores must match topk_indices")
+
+    sort_key = torch.where(valid_mask, topk_indices, torch.full_like(topk_indices, sk))
+    order = sort_key.argsort(dim=-1)
+    sorted_valid = torch.gather(valid_mask, dim=-1, index=order)
+    sorted_indices = torch.gather(topk_indices, dim=-1, index=order)
+    sorted_indices = sorted_indices.masked_fill(~sorted_valid, -1).contiguous()
+    if topk_scores is None:
+        return sorted_indices, None
+
+    sorted_scores = torch.gather(topk_scores, dim=-1, index=order)
+    sorted_scores = sorted_scores.masked_fill(~sorted_valid, invalid_score).contiguous()
+    return sorted_indices, sorted_scores
+
+
 def apply_starts_ends_mask_to_scores(
     scores: torch.Tensor, starts: torch.Tensor, ends: torch.Tensor, key_positions: torch.Tensor
 ) -> torch.Tensor:

megatron/core/transformer/experimental_attention_variant/dsa_tilelang_kernels.py

@@ -0,0 +1,142 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""TileLang backend hooks for optional fused DeepSeek sparse attention kernels."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional, Tuple
+
+import torch
+
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer.experimental_attention_variant.ops import tilelang_dsa
+
+if TYPE_CHECKING:
+    from megatron.core.packed_seq_params import PackedSeqParams
+    from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+def run_fused_qk_topk(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    weights: torch.Tensor,
+    index_topk: int,
+    starts: torch.Tensor,
+    ends: torch.Tensor,
+    block_size: int,
+    use_relu: bool = True,
+    use_local_indexer_varlen: bool = False,
+    single_packed_thd_sequence: bool = False,
+    local_packed_cp_rank: int = 0,
+    local_packed_cp_query_start: int = 0,
+    local_packed_cp_query_len: Optional[int] = None,
+    packed_seq_params: Optional[PackedSeqParams] = None,
+    cp_size: int = 1,
+) -> Optional[Tuple[torch.Tensor, Optional[torch.Tensor]]]:
+    """Adapt TileLang's indices-only result to the shared backend hook contract."""
+    topk_indices = tilelang_dsa.run_fused_qk_topk(
+        q,
+        k,
+        weights,
+        index_topk,
+        starts,
+        ends,
+        block_size,
+        use_relu,
+        use_local_indexer_varlen=use_local_indexer_varlen,
+        single_packed_thd_sequence=single_packed_thd_sequence,
+        local_packed_cp_rank=local_packed_cp_rank,
+        local_packed_cp_query_start=local_packed_cp_query_start,
+        local_packed_cp_query_len=local_packed_cp_query_len,
+        packed_seq_params=packed_seq_params,
+        cp_size=cp_size,
+    )
+    if topk_indices is None:
+        return None
+    return topk_indices, None
+
+
+def run_fused_qk_topk_with_loss(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    weights: torch.Tensor,
+    index_topk: int,
+    starts: torch.Tensor,
+    ends: torch.Tensor,
+    block_size: int,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    softmax_scale: float,
+    loss_coeff: float,
+    pg_collection: ProcessGroupCollection,
+    query_valid_rows: Optional[torch.Tensor] = None,
+    calculate_per_token_loss: bool = False,
+    use_relu: bool = True,
+    config: Optional["TransformerConfig"] = None,
+    use_local_indexer_varlen: bool = False,
+    single_packed_thd_sequence: bool = False,
+    local_packed_cp_rank: int = 0,
+    local_packed_cp_query_start: int = 0,
+    local_packed_cp_query_len: Optional[int] = None,
+    packed_seq_params: Optional[PackedSeqParams] = None,
+    cp_size: int = 1,
+) -> Optional[Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]]:
+    """Run fused TileLang indexer and sparse indexer loss."""
+    del config
+    result = tilelang_dsa.run_fused_qk_topk_with_loss(
+        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,
+        use_local_indexer_varlen=use_local_indexer_varlen,
+        single_packed_thd_sequence=single_packed_thd_sequence,
+        local_packed_cp_rank=local_packed_cp_rank,
+        local_packed_cp_query_start=local_packed_cp_query_start,
+        local_packed_cp_query_len=local_packed_cp_query_len,
+        packed_seq_params=packed_seq_params,
+        cp_size=cp_size,
+    )
+    if result is None:
+        return None
+    topk_indices, indexer_loss = result
+    return topk_indices, None, indexer_loss
+
+
+def run_fused_absorbed_sparse_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    topk_indices: torch.Tensor,
+    softmax_scale: float,
+    v_channels: int,
+    topk_length: Optional[torch.Tensor] = None,
+) -> Optional[torch.Tensor]:
+    """Run fused TileLang SparseMLA for absorbed DSA sparse attention."""
+    if topk_length is not None:
+        if topk_indices.ndim != 3 or topk_length.shape != topk_indices.shape[:-1]:
+            return None
+        positions = torch.arange(topk_indices.size(-1), device=topk_indices.device)
+        valid = positions < topk_length.to(dtype=torch.int64, device=topk_indices.device).unsqueeze(
+            -1
+        )
+        topk_indices = topk_indices.masked_fill(~valid, -1)
+    return tilelang_dsa.run_fused_absorbed_sparse_attention(
+        query, key, topk_indices, softmax_scale, v_channels
+    )
+
+
+__all__ = [
+    "run_fused_absorbed_sparse_attention",
+    "run_fused_qk_topk",
+    "run_fused_qk_topk_with_loss",
+]