diff --git a/megatron/core/transformer/experimental_attention_variant/dsa_cudnn_kernels.py b/megatron/core/transformer/experimental_attention_variant/dsa_cudnn_kernels.py index e26c46030aa..640db5c91ec 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa_cudnn_kernels.py +++ b/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]: diff --git a/megatron/core/transformer/experimental_attention_variant/dsa_layout.py b/megatron/core/transformer/experimental_attention_variant/dsa_layout.py index eec5a570a44..080ee2dc8ea 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa_layout.py +++ b/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: diff --git a/megatron/core/transformer/experimental_attention_variant/dsa_masking.py b/megatron/core/transformer/experimental_attention_variant/dsa_masking.py index 98126e63798..4c178b7dcea 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa_masking.py +++ b/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: diff --git a/megatron/core/transformer/experimental_attention_variant/dsa_tilelang_kernels.py b/megatron/core/transformer/experimental_attention_variant/dsa_tilelang_kernels.py new file mode 100644 index 00000000000..1d89e43e196 --- /dev/null +++ b/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", +] diff --git a/megatron/core/transformer/experimental_attention_variant/ops/indexer.py b/megatron/core/transformer/experimental_attention_variant/ops/indexer.py new file mode 100644 index 00000000000..2f59feb0776 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/ops/indexer.py @@ -0,0 +1,131 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import torch + +from .tilelang_indexer_bwd import HAVE_TILELANG as HAVE_TILELANG_INDEXER_BWD +from .tilelang_indexer_bwd import indexer_bwd_interface +from .tilelang_indexer_fwd import HAVE_TILELANG as HAVE_TILELANG_INDEXER_FWD +from .tilelang_indexer_fwd import indexer_fwd_interface + +HAVE_TILELANG_INDEXER = HAVE_TILELANG_INDEXER_BWD and HAVE_TILELANG_INDEXER_FWD + + +def pytorch_extract_topk_scores(logits, topk_indices, dim=-1): + """Gather top-k logits and mask invalid (-1) entries with -inf.""" + if logits.size(dim) == 0: + return torch.full( + topk_indices.shape, float("-inf"), dtype=logits.dtype, device=logits.device + ) + valid_mask = (topk_indices >= 0) & (topk_indices < logits.size(dim)) + safe_indices = topk_indices.clamp(min=0, max=logits.size(dim) - 1).to(torch.int64) + scores = torch.gather(logits, dim=dim, index=safe_indices) + scores = torch.where(valid_mask, scores, float("-inf")) + return scores + + +def _select_topk_from_logits( + logits: torch.Tensor, topk: int, mask_invalid: bool = True +) -> tuple[torch.Tensor, torch.Tensor]: + """Select top-k scores and int32 indices from indexer logits.""" + effective_topk = min(topk, logits.size(-1)) + if effective_topk > 0: + topk_scores, topk_indices = torch.topk(logits, effective_topk, dim=-1, sorted=False) + topk_indices = topk_indices.to(torch.int32) + if mask_invalid: + topk_indices = topk_indices.masked_fill(topk_scores == -torch.inf, -1) + return topk_scores, topk_indices + + empty_shape = logits.shape[:-1] + (0,) + topk_scores = torch.empty(empty_shape, dtype=logits.dtype, device=logits.device) + topk_indices = torch.empty(empty_shape, dtype=torch.int32, device=logits.device) + return topk_scores, topk_indices + + +class IndexerFunction(torch.autograd.Function): # pragma: no cover + """Autograd wrapper for fused tilelang indexer forward/backward.""" + + @staticmethod + def forward( + ctx, + index_q: torch.Tensor, + index_k: torch.Tensor, + weights: torch.Tensor, + cu_seqlen_ks: torch.Tensor, + cu_seqlen_ke: torch.Tensor, + topk: int, + topk_indices: torch.Tensor | None = None, + use_relu: bool = True, + ): + """Run fused indexer forward and optionally select top-k indices.""" + logits = indexer_fwd_interface( + index_q, + index_k, + weights, + cu_seqlen_ks, + cu_seqlen_ke, + clean_logits=True, + use_relu=use_relu, + ) + if topk_indices is None: + index_score, topk_indices = _select_topk_from_logits(logits, topk) + else: + index_score = pytorch_extract_topk_scores(logits, topk_indices) + + ctx.save_for_backward(index_q, index_k, weights, cu_seqlen_ks, cu_seqlen_ke, topk_indices) + ctx.use_relu = use_relu + return index_score, topk_indices + + @staticmethod + def backward(ctx, grad_scores, grad_indices): + """Propagate gradients through fused indexer outputs.""" + index_q, index_k, weights, cu_seqlen_ks, cu_seqlen_ke, topk_indices = ctx.saved_tensors + grad_q, grad_w, grad_k = indexer_bwd_interface( + index_q, weights, index_k, topk_indices, grad_scores, use_relu=ctx.use_relu + ) + return grad_q, grad_k, grad_w, None, None, None, None, None + + +def lighting_indexer( # pragma: no cover + index_q: torch.Tensor, + index_k: torch.Tensor, + weights: torch.Tensor, + cu_seqlen_ks: torch.Tensor, + cu_seqlen_ke: torch.Tensor, + topk: int, + topk_indices: torch.Tensor | None = None, + use_relu: bool = True, +): + """Compute indexer top-k scores/indices via the custom autograd function.""" + return IndexerFunction.apply( + index_q, index_k, weights, cu_seqlen_ks, cu_seqlen_ke, topk, topk_indices, use_relu + ) + + +def lighting_indexer_indices( # pragma: no cover + index_q: torch.Tensor, + index_k: torch.Tensor, + weights: torch.Tensor, + cu_seqlen_ks: torch.Tensor, + cu_seqlen_ke: torch.Tensor, + topk: int, + use_relu: bool = True, +): + """Compute TileLang indexer top-k indices without score/autograd bookkeeping.""" + with torch.no_grad(): + logits = indexer_fwd_interface( + index_q, + index_k, + weights, + cu_seqlen_ks, + cu_seqlen_ke, + clean_logits=True, + use_relu=use_relu, + ) + _, topk_indices = _select_topk_from_logits(logits, topk, mask_invalid=False) + return topk_indices + + +if not HAVE_TILELANG_INDEXER: + IndexerFunction = None + lighting_indexer = None + lighting_indexer_indices = None diff --git a/megatron/core/transformer/experimental_attention_variant/ops/sparse_mla.py b/megatron/core/transformer/experimental_attention_variant/ops/sparse_mla.py new file mode 100644 index 00000000000..76976ebe0aa --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/ops/sparse_mla.py @@ -0,0 +1,105 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import torch + +from .tilelang_sparse_mla_bwd import HAVE_TILELANG as HAVE_TILELANG_SPARSE_MLA_BWD +from .tilelang_sparse_mla_bwd import sparse_mla_bwd, sparse_mla_delta +from .tilelang_sparse_mla_fwd import HAVE_TILELANG as HAVE_TILELANG_SPARSE_MLA_FWD +from .tilelang_sparse_mla_fwd import sparse_mla_fwd_interface + +HAVE_TILELANG_SPARSE_MLA = HAVE_TILELANG_SPARSE_MLA_BWD and HAVE_TILELANG_SPARSE_MLA_FWD + + +def _canonicalize_batch_stride(tensor: torch.Tensor) -> torch.Tensor: + """Normalize a size-one batch stride without copying tensor data.""" + tensor = tensor.contiguous() + if tensor.ndim == 4 and tensor.size(0) == 1: + tensor = tensor.squeeze(0).unsqueeze(0) + return tensor + + +def _valid_head_mask(indices, num_heads): + valid_groups = indices.ge(0).any(dim=-1) + kv_group = valid_groups.size(-1) + if kv_group == num_heads: + return valid_groups + if num_heads % kv_group != 0: + raise RuntimeError( + f"SparseMLA heads must be divisible by kv_group, got heads={num_heads}, " + f"kv_group={kv_group}" + ) + return valid_groups.repeat_interleave(num_heads // kv_group, dim=-1) + + +def _zero_invalid_heads(tensor, valid_heads): + zero = torch.zeros((), dtype=tensor.dtype, device=tensor.device) + return torch.where(valid_heads.unsqueeze(-1), tensor, zero) + + +class SparseMLA(torch.autograd.Function): # pragma: no cover + """Autograd wrapper around tilelang sparse-MLA forward/backward kernels.""" + + @staticmethod + def forward(ctx, q, kv, indices, scaling): + """ + Args: + q: Query tensor (seq_len, heads, dim_plus_tail_dim) or + (batch, seq_len, heads, dim_plus_tail_dim) + kv: Key-Value tensor (seq_len_kv, kv_group, dim_plus_tail_dim) or + (batch, seq_len_kv, kv_group, dim_plus_tail_dim) + indices: Sparse indices tensor (seq_len, kv_group, topk) or + (batch, seq_len, kv_group, topk) + + Returns: + out: Output tensor (seq_len, heads, dim) or (batch, seq_len, heads, dim) + """ + indices = _canonicalize_batch_stride(indices) + q = _canonicalize_batch_stride(q) + kv = _canonicalize_batch_stride(kv) + ctx.scaling = scaling + valid_heads = _valid_head_mask(indices, q.size(-2)) + tl_out, tl_lse = sparse_mla_fwd_interface(q, kv, indices, sm_scale=scaling) + tl_out = _zero_invalid_heads(tl_out, valid_heads) + lse_zero = torch.zeros((), dtype=tl_lse.dtype, device=tl_lse.device) + tl_lse = torch.where(valid_heads, tl_lse, lse_zero) + + # Do not save tl_out/tl_lse: backward recomputes them just long enough to form + # delta and run the kernel. Saved inputs still go through autograd's saved-tensor + # hooks/offload path and retain_graph can recompute these tensors again. + ctx.save_for_backward(q, kv, indices, valid_heads) + + return tl_out, tl_lse + + @staticmethod + def backward(ctx, grad_output, grad_lse): + """ + Args: + grad_output: Gradient of the loss with respect to output + + Returns: + Gradients for q, kv, and indices (None for indices) + """ + q, kv, indices, valid_heads = ctx.saved_tensors + scaling = ctx.scaling + grad_output = grad_output.contiguous() + grad_output = _zero_invalid_heads(grad_output, valid_heads) + with torch.no_grad(): + tl_out, tl_lse = sparse_mla_fwd_interface(q, kv, indices, sm_scale=scaling) + tl_out = _zero_invalid_heads(tl_out, valid_heads) + lse_zero = torch.zeros((), dtype=tl_lse.dtype, device=tl_lse.device) + tl_lse = torch.where(valid_heads, tl_lse, lse_zero) + delta = sparse_mla_delta(tl_out, grad_output) + del tl_out + + tl_dq, tl_dkv = sparse_mla_bwd( + q, kv, None, grad_output, indices, tl_lse, sm_scale=scaling, delta=delta + ) + tl_dq = _zero_invalid_heads(tl_dq, valid_heads) + del tl_lse + + # Return gradients for each input (None for indices as it's not differentiable) + return tl_dq, tl_dkv, None, None + + +if not HAVE_TILELANG_SPARSE_MLA: + SparseMLA = None diff --git a/megatron/core/transformer/experimental_attention_variant/ops/tilelang_dsa.py b/megatron/core/transformer/experimental_attention_variant/ops/tilelang_dsa.py new file mode 100644 index 00000000000..bedcdb4ca1f --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/ops/tilelang_dsa.py @@ -0,0 +1,908 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""TileLang-backed DSA hook implementations. + +This module keeps TileLang-specific batching, chunking, and sparse-KL streaming +out of the backend-neutral DSA control flow in ``dsa.py``. +""" + +from collections import OrderedDict +from typing import TYPE_CHECKING, Optional, Tuple + +import torch + +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.experimental_attention_variant import ( + dsa_indexer_loss, + dsa_layout, + dsa_masking, +) +from megatron.core.utils import get_pg_size + +if TYPE_CHECKING: + from megatron.core.packed_seq_params import PackedSeqParams + +try: + from megatron.core.transformer.experimental_attention_variant.ops.indexer import ( + lighting_indexer, + lighting_indexer_indices, + ) + from megatron.core.transformer.experimental_attention_variant.ops.tilelang_indexer_bwd import ( + is_supported_indexer_bwd_head_count, + ) +except (ImportError, OSError): + is_supported_indexer_bwd_head_count = None + lighting_indexer = None + lighting_indexer_indices = None + +try: + from megatron.core.transformer.experimental_attention_variant.ops.sparse_mla import SparseMLA +except (ImportError, OSError): + SparseMLA = None + +try: + from megatron.core.transformer.experimental_attention_variant.ops.tilelang_indexer_loss import ( + SparseIndexerKLLoss, + sparse_indexer_target_interface, + ) +except (ImportError, OSError): + SparseIndexerKLLoss = None + sparse_indexer_target_interface = None + + +# Reusable no-grad scratch buffers keyed by (name, shape, dtype, device). +_DSA_SCRATCH_CACHE_MAX_ENTRIES = 128 +_DSA_SCRATCH_CACHE_MAX_BYTES = 512 * 1024 * 1024 +_DSA_SCRATCH_CACHE = OrderedDict() +_DSA_SCRATCH_CACHE_TOTAL_BYTES = 0 + + +def _scratch_buffer_bytes(buf: torch.Tensor) -> int: + return buf.numel() * buf.element_size() + + +def _is_supported_sparse_mla_head_count(heads: int, kv_group: int = 1) -> bool: + """Return whether TileLang SparseMLA supports this query/KV head grouping. + + The forward and backward kernels pad ``head_kv`` to ``max(next_power_of_2(head_kv), 16)`` + and index the unpadded head dimension by that padded count with no head-dim bound, so they + only stay in bounds when no padding occurs. That requires ``head_kv`` (= ``heads // + kv_group``) to be a power of two and at least 16; any other value (e.g. 48, 192, or < 16) + must fall back to the unfused path rather than read/write past the real head count. + """ + if kv_group <= 0 or heads % kv_group != 0: + return False + head_kv = heads // kv_group + return head_kv >= 16 and (head_kv & (head_kv - 1)) == 0 + + +def _all_bfloat16(*tensors: torch.Tensor) -> bool: + return all(tensor.dtype == torch.bfloat16 for tensor in tensors) + + +def _evict_scratch_cache_if_needed() -> None: + """Bound scratch cache growth by LRU eviction.""" + global _DSA_SCRATCH_CACHE_TOTAL_BYTES + while ( + len(_DSA_SCRATCH_CACHE) > _DSA_SCRATCH_CACHE_MAX_ENTRIES + or _DSA_SCRATCH_CACHE_TOTAL_BYTES > _DSA_SCRATCH_CACHE_MAX_BYTES + ): + _, buf = _DSA_SCRATCH_CACHE.popitem(last=False) + _DSA_SCRATCH_CACHE_TOTAL_BYTES -= _scratch_buffer_bytes(buf) + + +def _get_scratch_buffer( + name: str, shape: Tuple[int, ...], dtype: torch.dtype, device: torch.device +) -> torch.Tensor: + """Get a reusable scratch tensor for temporary no-grad workspaces.""" + global _DSA_SCRATCH_CACHE_TOTAL_BYTES + key = (name, shape, dtype, device) + buf = _DSA_SCRATCH_CACHE.pop(key, None) + if buf is not None: + _DSA_SCRATCH_CACHE_TOTAL_BYTES -= _scratch_buffer_bytes(buf) + else: + buf = torch.empty(shape, dtype=dtype, device=device) + _DSA_SCRATCH_CACHE[key] = buf + _DSA_SCRATCH_CACHE_TOTAL_BYTES += _scratch_buffer_bytes(buf) + _evict_scratch_cache_if_needed() + return buf + + +def _topk_valid_mask( + topk_indices: torch.Tensor, starts: torch.Tensor, ends: torch.Tensor +) -> torch.Tensor: + """Compute the row-wise [start, end) validity mask for fused indexer outputs.""" + starts_for_cmp = starts.to(device=topk_indices.device, dtype=topk_indices.dtype).unsqueeze(-1) + ends_for_cmp = ends.to(device=topk_indices.device, dtype=topk_indices.dtype).unsqueeze(-1) + return (topk_indices >= starts_for_cmp) & (topk_indices < ends_for_cmp) + + +def _sanitize_fused_topk_indices( + topk_indices: torch.Tensor, starts: torch.Tensor, ends: torch.Tensor +) -> torch.Tensor: + """Mask fused indexer outputs in place and return the validity mask.""" + valid = _topk_valid_mask(topk_indices, starts, ends) + topk_indices.masked_fill_(~valid, -1) + return valid + + +def _sanitize_fused_topk_outputs( + topk_indices: torch.Tensor, + starts: torch.Tensor, + ends: torch.Tensor, + topk_scores: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Mask fused indexer outputs and optional scores to row-wise key bounds.""" + valid = _topk_valid_mask(topk_indices, starts, ends) + sanitized_indices = topk_indices.masked_fill(~valid, -1) + if topk_scores is not None: + topk_scores = topk_scores.masked_fill(~valid, float("-inf")) + return sanitized_indices, topk_scores + + +def _build_packed_cp_indexer_inputs( + index_k: torch.Tensor, + starts: torch.Tensor, + ends: torch.Tensor, + *, + packed_seq_params: "PackedSeqParams", + cp_size: int, + cp_rank: int, + single_packed_thd_sequence: bool, + local_query_start: int, + local_query_len: int, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Pack CP front/back key prefixes and translate query bounds to the packed key space.""" + if cp_size <= 1 or not 0 <= cp_rank < cp_size: + raise RuntimeError("packed CP TileLang indexer requires a valid CP rank and cp_size > 1") + if local_query_start < 0 or local_query_start + starts.numel() > local_query_len: + raise RuntimeError( + "packed CP TileLang indexer received an invalid local query slice: " + f"start={local_query_start}, rows={starts.numel()}, local_rows={local_query_len}" + ) + + cu_q, cu_k = dsa_layout.get_packed_qk_cu_seqlens(packed_seq_params) + device = index_k.device + sk = index_k.size(0) + layout = dsa_layout.build_packed_cp_indexer_layout( + cu_q.to(device=device), + cu_k.to(device=device), + cp_size=cp_size, + cp_rank=cp_rank, + key_size=sk, + local_key_layout=single_packed_thd_sequence and sk == local_query_len, + ) + segmented_k = index_k.index_select(0, layout.source_indices).contiguous() + + segment_ids_q = torch.repeat_interleave( + torch.arange(layout.segment_q_lengths.numel(), device=device), + layout.segment_q_lengths, + output_size=local_query_len, + ) + row_start = local_query_start + row_end = row_start + starts.numel() + row_segment_ids = segment_ids_q[row_start:row_end] + row_segment_starts = layout.segment_cu_k[:-1].index_select(0, row_segment_ids) + row_segment_ends = row_segment_starts + layout.segment_k_lengths.index_select( + 0, row_segment_ids + ) + row_global_starts = layout.segment_key_starts.index_select(0, row_segment_ids) + + local_starts = row_segment_starts + starts.to(torch.int64) - row_global_starts + local_ends = row_segment_starts + ends.to(torch.int64) - row_global_starts + local_starts = torch.maximum(local_starts, row_segment_starts) + local_starts = torch.minimum(local_starts, row_segment_ends) + local_ends = torch.maximum(local_ends, local_starts) + local_ends = torch.minimum(local_ends, row_segment_ends) + return ( + segmented_k, + local_starts.to(torch.int32).contiguous(), + local_ends.to(torch.int32).contiguous(), + layout.source_indices, + ) + + +def _remap_segmented_topk_indices( + topk_indices: torch.Tensor, source_indices: torch.Tensor +) -> torch.Tensor: + """Map valid indices from a segmented key tensor back to the original packed key tensor.""" + valid = topk_indices >= 0 + safe_indices = topk_indices.clamp(min=0).reshape(-1).to(torch.int64) + global_indices = source_indices.index_select(0, safe_indices).view_as(topk_indices) + return torch.where(valid, global_indices.to(topk_indices.dtype), topk_indices) + + +def fused_qk_topk_lighting( + 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[torch.Tensor]: + """Run fused TileLang indexer and return top-k indices [b, sq, topk].""" + if lighting_indexer_indices is None: + return None + if q.ndim != 4 or k.ndim != 3 or weights.ndim != 3: + return None + if not _all_bfloat16(q, k): + return None + + sq, b = q.size(0), q.size(1) + if k.size(1) != b or weights.size(1) != b: + return None + starts = starts.contiguous() + ends = ends.contiguous() + + topk_k = min(index_topk, k.size(0)) + topk_out = torch.empty((b, sq, topk_k), dtype=torch.int32, device=q.device) + for bi in range(b): + index_q = q[:, bi].contiguous() + index_k = k[:, bi].contiguous() + index_w = weights[:, bi].float().contiguous() + local_starts = starts + local_ends = ends + source_indices = None + if b == 1 and use_local_indexer_varlen and packed_seq_params is not None and cp_size > 1: + local_query_len = ( + local_packed_cp_query_len if local_packed_cp_query_len is not None else sq + ) + index_k, local_starts, local_ends, source_indices = _build_packed_cp_indexer_inputs( + index_k, + starts, + ends, + packed_seq_params=packed_seq_params, + cp_size=cp_size, + cp_rank=local_packed_cp_rank, + single_packed_thd_sequence=single_packed_thd_sequence, + local_query_start=local_packed_cp_query_start, + local_query_len=local_query_len, + ) + for start in range(0, sq, block_size): + end = min(start + block_size, sq) + topk_indices = lighting_indexer_indices( + index_q[start:end], + index_k, + index_w[start:end], + local_starts[start:end], + local_ends[start:end], + topk_k, + use_relu=use_relu, + ) + _sanitize_fused_topk_indices( + topk_indices, starts=local_starts[start:end], ends=local_ends[start:end] + ) + if source_indices is not None: + topk_indices = _remap_segmented_topk_indices(topk_indices, source_indices) + topk_out[bi, start:end].copy_(topk_indices) + + return topk_out + + +@torch.no_grad() +def _compute_topk_target_chunk_sum( + *, + query_h: torch.Tensor, + key_shared: Optional[torch.Tensor], + key_per_head: Optional[torch.Tensor], + s0: int, + s1: int, + idx_seq: torch.Tensor, + valid_seq: torch.Tensor, + softmax_scale: float, + head_chunk_size: int, + topk_chunk_size: int, + sk: int, + hn: int, +) -> torch.Tensor: + """Compute unnormalized target probability mass on top-k support for one sequence chunk.""" + s_len = s1 - s0 + topk = idx_seq.size(-1) + device = query_h.device + np = query_h.size(0) + + attn_chunk_sum = _get_scratch_buffer("kl_attn_chunk_sum", (s_len, topk), torch.float32, device) + attn_chunk_sum.zero_() + + for h0 in range(0, np, head_chunk_size): + h1 = min(h0 + head_chunk_size, np) + h_chunk = h1 - h0 + q_chunk = query_h[h0:h1, s0:s1, :] + q_chunk_float = q_chunk.float() + + if key_shared is None: + key_chunk = key_per_head[h0:h1] + flat_keys = key_chunk.reshape(h_chunk * sk, hn) + head_offsets = ( + torch.arange(h_chunk, device=device, dtype=torch.int64).view(-1, 1, 1) * sk + ) + else: + flat_keys = None + head_offsets = None + + # Two-pass online softmax over top-k chunks: + # 1) compute row-wise max and denominator; 2) recompute and accumulate probabilities. + # These accumulators are rebound to fresh tensors each top-k chunk, so they + # cannot reuse a scratch buffer in place. + running_max = torch.full( + (h_chunk, s_len), float("-inf"), dtype=torch.float32, device=device + ) + running_denom = torch.zeros((h_chunk, s_len), dtype=torch.float32, device=device) + + def _chunk_logits(idx_topk, valid_topk_chunk, k_len): + if key_shared is not None: + key_sel = key_shared.index_select(0, idx_topk.reshape(-1)).view(s_len, k_len, hn) + logits = torch.einsum("hsd,skd->hsk", q_chunk_float, key_sel.float()) + else: + flat_idx = idx_topk.unsqueeze(0) + head_offsets + key_sel = flat_keys.index_select(0, flat_idx.reshape(-1)).view( + h_chunk, s_len, k_len, hn + ) + logits = (q_chunk_float.unsqueeze(2) * key_sel.float()).sum(dim=-1) + logits = logits * softmax_scale + return logits.masked_fill(~valid_topk_chunk.unsqueeze(0), float("-inf")) + + for t0 in range(0, topk, topk_chunk_size): + t1 = min(t0 + topk_chunk_size, topk) + logits = _chunk_logits(idx_seq[:, t0:t1], valid_seq[:, t0:t1], t1 - t0) + chunk_max = logits.max(dim=-1).values + new_running_max = torch.maximum(running_max, chunk_max) + max_for_exp = torch.where( + torch.isfinite(new_running_max), new_running_max, torch.zeros_like(new_running_max) + ) + alpha = torch.exp(running_max - max_for_exp) + p_chunk = torch.exp(logits - max_for_exp.unsqueeze(-1)) + running_denom = running_denom * alpha + p_chunk.sum(dim=-1) + running_max = new_running_max + + stable_max = torch.where( + torch.isfinite(running_max), running_max, torch.zeros_like(running_max) + ) + inverse_denom = running_denom.clamp_min(1e-10).reciprocal() + for t0 in range(0, topk, topk_chunk_size): + t1 = min(t0 + topk_chunk_size, topk) + logits = _chunk_logits(idx_seq[:, t0:t1], valid_seq[:, t0:t1], t1 - t0) + probs = torch.exp(logits - stable_max.unsqueeze(-1)) * inverse_denom.unsqueeze(-1) + attn_chunk_sum[:, t0:t1] += probs.sum(dim=0) + + return attn_chunk_sum + + +def _compute_sparse_topk_kl_chunk( + target_chunk: torch.Tensor, index_logits_chunk: torch.Tensor, valid_seq: torch.Tensor +) -> torch.Tensor: + """Compute KL(target || index) sum for one [s_chunk, topk] chunk.""" + index_logits_chunk = index_logits_chunk.to(dtype=torch.float32, device=target_chunk.device) + target_chunk = target_chunk.to(dtype=torch.float32, device=index_logits_chunk.device) + with torch.no_grad(): + index_log_scores_chunk = dsa_masking.masked_log_softmax( + index_logits_chunk.detach(), valid_seq, dim=-1 + ) + index_scores_chunk = index_log_scores_chunk.exp().masked_fill(~valid_seq, 0.0) + kl_value = dsa_indexer_loss.indexer_kl_sum(target_chunk, index_log_scores_chunk, valid_seq) + grad_logits = (index_scores_chunk - target_chunk).masked_fill(~valid_seq, 0.0) + index_logits_for_grad = index_logits_chunk.masked_fill(~valid_seq, 0.0) + grad_surrogate = (index_logits_for_grad * grad_logits).sum() + return grad_surrogate + (kl_value - grad_surrogate).detach() + + +def _can_use_fused_sparse_indexer_target( + query: torch.Tensor, key: Optional[torch.Tensor], topk_indices: torch.Tensor +) -> bool: + """Return whether the fused TileLang target kernel supports these tensors.""" + return ( + sparse_indexer_target_interface is not None + and key is not None + and query.is_cuda + and key.is_cuda + and topk_indices.is_cuda + and query.ndim == 3 + and key.ndim == 2 + and query.dtype == torch.bfloat16 + and key.dtype == torch.bfloat16 + and query.size(-1) == key.size(-1) + and query.size(-1) % 16 == 0 + and topk_indices.ndim == 2 + and topk_indices.size(-1) % 64 == 0 + ) + + +def _can_use_fused_sparse_indexer_kl( + target: torch.Tensor, index_logits: torch.Tensor, valid_mask: torch.Tensor +) -> bool: + """Return whether the fused TileLang KL/score-gradient kernel supports these tensors.""" + return ( + SparseIndexerKLLoss is not None + and target.is_cuda + and index_logits.is_cuda + and valid_mask.is_cuda + and target.dtype == torch.float32 + and index_logits.dtype == torch.float32 + and valid_mask.dtype == torch.bool + and target.shape == index_logits.shape == valid_mask.shape + and target.ndim == 2 + and target.size(-1) % 256 == 0 + ) + + +def _canonicalize_topk_scores_for_tp_reduce( + topk_indices: torch.Tensor, topk_scores: torch.Tensor, *, sk: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """Sort selected top-k slots by key index before slot-wise TP reductions.""" + valid = topk_indices >= 0 + topk_indices, topk_scores = dsa_masking.sort_topk_by_index( + topk_indices, valid, sk=sk, topk_scores=topk_scores + ) + assert topk_scores is not None + return topk_indices, topk_scores + + +def _accumulate_topk_kl_chunk( + *, + target_chunk: torch.Tensor, + index_logits_chunk: torch.Tensor, + valid_seq: torch.Tensor, + kl_sum: torch.Tensor, +) -> torch.Tensor: + """Normalize one target chunk and accumulate its sparse KL contribution.""" + if _can_use_fused_sparse_indexer_kl(target_chunk, index_logits_chunk, valid_seq): + return kl_sum + SparseIndexerKLLoss.apply( + target_chunk.contiguous(), index_logits_chunk.contiguous(), valid_seq.contiguous() + ) + normalized_target = dsa_indexer_loss.normalize_indexer_target_(target_chunk) + return kl_sum + _compute_sparse_topk_kl_chunk( + target_chunk=normalized_target, index_logits_chunk=index_logits_chunk, valid_seq=valid_seq + ) + + +def _stage_topk_target_chunk( + target_chunk: torch.Tensor, + *, + slot_prefix: str, + slot: int, + device: torch.device, + tp_group: torch.distributed.ProcessGroup, + tp_size: int, +) -> Tuple[torch.Tensor, Optional[torch.distributed.Work]]: + """Copy chunk into scratch slot and optionally launch async TP all-reduce.""" + target_chunk_work = _get_scratch_buffer( + f"{slot_prefix}_slot{slot}", tuple(target_chunk.shape), torch.float32, device + ) + target_chunk_work.copy_(target_chunk) + if tp_size > 1: + handle = torch.distributed.all_reduce(target_chunk_work, group=tp_group, async_op=True) + else: + handle = None + return target_chunk_work, handle + + +def _consume_pending_topk_kl_chunk( + *, + pending_handle: Optional[torch.distributed.Work], + pending_target_chunk: Optional[torch.Tensor], + pending_index_logits: Optional[torch.Tensor], + pending_valid_seq: Optional[torch.Tensor], + kl_sum: torch.Tensor, +) -> torch.Tensor: + """Finalize one pending chunk and accumulate its KL contribution into ``kl_sum``.""" + if pending_target_chunk is None: + return kl_sum + if pending_handle is not None: + pending_handle.wait() + return _accumulate_topk_kl_chunk( + target_chunk=pending_target_chunk, + index_logits_chunk=pending_index_logits, + valid_seq=pending_valid_seq, + kl_sum=kl_sum, + ) + + +def _enqueue_topk_kl_chunk( + *, + target_chunk: torch.Tensor, + index_logits_chunk: torch.Tensor, + valid_seq: torch.Tensor, + slot_prefix: str, + chunk_id: int, + device: torch.device, + tp_group: torch.distributed.ProcessGroup, + tp_size: int, + pending_handle: Optional[torch.distributed.Work], + pending_target_chunk: Optional[torch.Tensor], + pending_index_logits: Optional[torch.Tensor], + pending_valid_seq: Optional[torch.Tensor], + kl_sum: torch.Tensor, +) -> Tuple[ + torch.Tensor, + int, + Optional[torch.distributed.Work], + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[torch.Tensor], +]: + """Stage a new KL chunk, consume previous pending chunk, and update pending state.""" + slot = chunk_id & 1 + target_chunk_work, current_handle = _stage_topk_target_chunk( + target_chunk, + slot_prefix=slot_prefix, + slot=slot, + device=device, + tp_group=tp_group, + tp_size=tp_size, + ) + kl_sum = _consume_pending_topk_kl_chunk( + pending_handle=pending_handle, + pending_target_chunk=pending_target_chunk, + pending_index_logits=pending_index_logits, + pending_valid_seq=pending_valid_seq, + kl_sum=kl_sum, + ) + return (kl_sum, chunk_id + 1, current_handle, target_chunk_work, index_logits_chunk, valid_seq) + + +def fused_qk_topk_lighting_with_streaming_sparse_kl( + 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, + seq_chunk_size: int = 512, + head_chunk_size: int = 16, + topk_chunk_size: int = 1024, + 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, torch.Tensor]]: + """Run the fused TileLang indexer with streaming sparse KL accumulation. + + The objective matches ``compute_dsa_indexer_loss`` on the selected top-k support. TileLang + streams query/head/top-k chunks and overlaps TP target reduction to avoid materializing dense + scores; its custom gradient surrogate supplies the same log-softmax gradient for fused indexer + logits. Target normalization, KL evaluation, and token reduction use the shared backend-neutral + helpers in ``dsa_indexer_loss``. + """ + if lighting_indexer is None: + return None + if q.ndim != 4 or k.ndim != 3 or weights.ndim != 3: + return None + if not _all_bfloat16(q, k): + return None + if is_supported_indexer_bwd_head_count is None or not is_supported_indexer_bwd_head_count( + q.size(2) + ): + return None + + query, _ = dsa_layout.ensure_sbhd(query, "query") + key, _ = dsa_layout.ensure_sbhd(key, "key") + sq, b = q.size(0), q.size(1) + sq_q, b_q, np, hn = query.size() + sk, b_k, nk, hk = key.size() + if k.size(1) != b or weights.size(1) != b: + return None + if sq_q != sq or b_q != b or b_k != b or hk != hn: + return None + if nk != 1 and nk != np: + return None + query_valid_rows = dsa_masking.normalize_query_valid_rows( + query_valid_rows, b=b, sq=sq, device=query.device + ) + + starts = starts.contiguous() + ends = ends.contiguous() + + topk_out = None + kl_sum = torch.zeros((), dtype=torch.float32, device=q.device) + tp_size = get_pg_size(pg_collection.tp) + pending_handle = None + pending_target_chunk = None + pending_index_logits = None + pending_valid_seq = None + chunk_id = 0 + for bi in range(b): + query_h = query[:, bi].permute(1, 0, 2).contiguous() + if nk == 1: + key_shared = key[:, bi, 0].contiguous() + key_per_head = None + else: + key_shared = None + key_per_head = key[:, bi].permute(1, 0, 2).contiguous() + + index_q = q[:, bi].contiguous() + index_k = k[:, bi].contiguous() + index_w = weights[:, bi].float().contiguous() + local_starts = starts + local_ends = ends + source_indices = None + if b == 1 and use_local_indexer_varlen and packed_seq_params is not None and cp_size > 1: + local_query_len = ( + local_packed_cp_query_len if local_packed_cp_query_len is not None else sq + ) + index_k, local_starts, local_ends, source_indices = _build_packed_cp_indexer_inputs( + index_k, + starts, + ends, + packed_seq_params=packed_seq_params, + cp_size=cp_size, + cp_rank=local_packed_cp_rank, + single_packed_thd_sequence=single_packed_thd_sequence, + local_query_start=local_packed_cp_query_start, + local_query_len=local_query_len, + ) + + for start in range(0, sq, block_size): + end = min(start + block_size, sq) + topk_scores, topk_indices = lighting_indexer( + index_q[start:end], + index_k, + index_w[start:end], + local_starts[start:end], + local_ends[start:end], + min(index_topk, k.size(0)), + topk_indices=None, + use_relu=use_relu, + ) + topk_indices, topk_scores = _sanitize_fused_topk_outputs( + topk_indices=topk_indices, + starts=local_starts[start:end], + ends=local_ends[start:end], + topk_scores=topk_scores, + ) + if source_indices is not None: + topk_indices = _remap_segmented_topk_indices(topk_indices, source_indices) + if tp_size > 1: + topk_indices, topk_scores = _canonicalize_topk_scores_for_tp_reduce( + topk_indices, topk_scores, sk=sk + ) + + if topk_out is None: + topk_out = torch.empty( + (b, sq, topk_indices.size(-1)), + dtype=topk_indices.dtype, + device=topk_indices.device, + ) + topk_out[bi, start:end].copy_(topk_indices) + + s_len = end - start + for rel_start in range(0, s_len, seq_chunk_size): + rel_end = min(rel_start + seq_chunk_size, s_len) + abs_start = start + rel_start + abs_end = start + rel_end + + idx_seq_raw = topk_indices[rel_start:rel_end].to(device=query.device) + valid_seq = idx_seq_raw >= 0 + if query_valid_rows is not None: + row_valid = query_valid_rows[bi, abs_start:abs_end] + valid_seq = valid_seq & row_valid.unsqueeze(-1) + loss_topk_indices = idx_seq_raw.masked_fill(~valid_seq, -1).contiguous() + query_chunk = query[abs_start:abs_end, bi].contiguous() + if _can_use_fused_sparse_indexer_target(query_chunk, key_shared, loss_topk_indices): + target_chunk = sparse_indexer_target_interface( + query_chunk, key_shared, loss_topk_indices, softmax_scale + ) + else: + target_chunk = _compute_topk_target_chunk_sum( + query_h=query_h, + key_shared=key_shared, + key_per_head=key_per_head, + s0=abs_start, + s1=abs_end, + idx_seq=idx_seq_raw.clamp(min=0).to(torch.int64), + valid_seq=valid_seq, + softmax_scale=softmax_scale, + head_chunk_size=head_chunk_size, + topk_chunk_size=topk_chunk_size, + sk=sk, + hn=hn, + ) + index_logits_chunk = topk_scores[rel_start:rel_end] + ( + kl_sum, + chunk_id, + pending_handle, + pending_target_chunk, + pending_index_logits, + pending_valid_seq, + ) = _enqueue_topk_kl_chunk( + target_chunk=target_chunk, + index_logits_chunk=index_logits_chunk, + valid_seq=valid_seq, + slot_prefix="stream_kl_target", + chunk_id=chunk_id, + device=query.device, + tp_group=pg_collection.tp, + tp_size=tp_size, + pending_handle=pending_handle, + pending_target_chunk=pending_target_chunk, + pending_index_logits=pending_index_logits, + pending_valid_seq=pending_valid_seq, + kl_sum=kl_sum, + ) + kl_sum = _consume_pending_topk_kl_chunk( + pending_handle=pending_handle, + pending_target_chunk=pending_target_chunk, + pending_index_logits=pending_index_logits, + pending_valid_seq=pending_valid_seq, + kl_sum=kl_sum, + ) + + if topk_out is None: + return None + valid_row_count = query_valid_rows.sum() if query_valid_rows is not None else None + kl_div = dsa_indexer_loss.reduce_indexer_kl_sum( + kl_sum, + num_rows=b * sq, + calculate_per_token_loss=calculate_per_token_loss, + valid_row_count=valid_row_count, + ) + return topk_out, kl_div * loss_coeff + + +def fused_sparse_mla_absorbed( + query: torch.Tensor, + key: torch.Tensor, + topk_indices: torch.Tensor, + softmax_scale: float, + v_channels: int, +) -> Optional[torch.Tensor]: + """Run fused SparseMLA kernel for absorbed-MLA path.""" + if SparseMLA is None: + return None + + if query.ndim != 4 or key.ndim != 4 or topk_indices.ndim != 3: + return None + if not _all_bfloat16(query, key): + return None + if key.size(2) != 1: + return None + if query.size(1) != key.size(1) or topk_indices.size(0) != query.size(1): + return None + if topk_indices.size(1) != query.size(0): + return None + if query.size(-1) != key.size(-1): + return None + if query.size(-1) != 576 or v_channels != 512: + # Current copied TileLang kernels are specialized for GLM5/DeepSeek V3.2 absorbed dims. + return None + query_heads = query.size(2) + if query_heads <= 0: + return None + kernel_heads = max(query_heads, 16) + if not _is_supported_sparse_mla_head_count(kernel_heads, kv_group=key.size(2)): + return None + if topk_indices.size(-1) % 64 != 0: + return None + + query_bshd = query.permute(1, 0, 2, 3).contiguous() + if kernel_heads != query_heads: + # SparseMLA uses a minimum 16-head tile without head bounds. Pad the caller + # tensor so small TP shards stay in bounds, then discard those heads below. + query_bshd = torch.nn.functional.pad(query_bshd, (0, 0, 0, kernel_heads - query_heads)) + key_bshd = key.permute(1, 0, 2, 3).contiguous() + indices_bsgk = topk_indices.unsqueeze(2).to(torch.int32).contiguous() + out, _ = SparseMLA.apply(query_bshd, key_bshd, indices_bsgk, softmax_scale) + if out.ndim != 4 or out.size(2) != kernel_heads or out.size(-1) != v_channels: + return None + out = out[:, :, :query_heads] + return out.permute(1, 0, 2, 3).contiguous() + + +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[torch.Tensor]: + """Optional fused indexer hook backed by TileLang.""" + return fused_qk_topk_lighting( + 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, + ) + + +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, + 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, torch.Tensor]]: + """Optional fused indexer+loss hook backed by TileLang.""" + return fused_qk_topk_lighting_with_streaming_sparse_kl( + 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, + ) + + +def run_fused_absorbed_sparse_attention( + query: torch.Tensor, + key: torch.Tensor, + topk_indices: torch.Tensor, + softmax_scale: float, + v_channels: int, +) -> Optional[torch.Tensor]: + """Optional fused sparse-attention hook backed by TileLang.""" + return fused_sparse_mla_absorbed(query, key, topk_indices, softmax_scale, v_channels) diff --git a/megatron/core/transformer/experimental_attention_variant/ops/tilelang_indexer_bwd.py b/megatron/core/transformer/experimental_attention_variant/ops/tilelang_indexer_bwd.py new file mode 100644 index 00000000000..33c5627728b --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/ops/tilelang_indexer_bwd.py @@ -0,0 +1,234 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# ruff: noqa +# Adapted from: +# https://github.com/tile-ai/tilelang/blob/4956b5835fa554af6c03d4a6289cad44bf310869/ +# examples/dsa_sparse_finetune/indexer_bwd.py +import threading +from collections import OrderedDict + +import torch + +from .tilelang_utils import ( + HAVE_TILELANG, + T, + _get_cached_kernel, + _next_power_of_two, + _round_up, + require_tilelang, +) +from .tilelang_utils import tilelang as tl +from .tilelang_utils import tilelang_jit + +BF16 = T.bfloat16 if HAVE_TILELANG else None +FP32 = T.float32 if HAVE_TILELANG else None +INT32 = T.int32 if HAVE_TILELANG else None +_tilelang_indexer_bwd_kernel_cache = OrderedDict() +_tilelang_indexer_bwd_cache_lock = threading.Lock() + +if HAVE_TILELANG: + pass_configs = { + tl.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tl.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + } +else: + pass_configs = {} + + +def _canonical_topk(topk: int, block_i: int = 32) -> int: + return _round_up(_next_power_of_two(topk), block_i) + + +def is_supported_indexer_bwd_head_count(heads: int) -> bool: + """Return whether TileLang indexer backward supports this indexer head count.""" + return heads <= 64 and heads % 8 == 0 + + +def _get_indexer_bwd_kernel(heads: int, dim: int, topk: int, use_relu: bool = True): + num_threads = 32 if heads < 16 else 128 + return _get_cached_kernel( + _tilelang_indexer_bwd_kernel_cache, + _tilelang_indexer_bwd_cache_lock, + (heads, dim, topk, use_relu), + lambda: tl_indexer_bwd_impl(heads, dim, topk, num_threads=num_threads, use_relu=use_relu), + ) + + +@tilelang_jit(pass_configs=pass_configs) +def tl_indexer_bwd_impl( # pragma: no cover + heads: int, + dim: int, + topk: int, + block_I: int = 32, + num_stages: int = 0, + num_threads: int = 128, + use_relu: bool = True, +): + """Build tilelang backward kernel for sparse indexer.""" + require_tilelang() + assert num_stages == 0 + assert topk == tl.math.next_power_of_2(topk) + assert topk % block_I == 0 + assert heads <= 64 and heads % 8 == 0 + seq_len = T.symbolic("seq_len") + q_seq_len = T.symbolic("q_seq_len") + + dtype: str = BF16 + accum_dtype: str = FP32 + index_q_shape = [q_seq_len, heads, dim] + weights_shape = [q_seq_len, heads] + index_k_shape = [seq_len, dim] + shape_p = [q_seq_len, topk] + topk_indices_shape = [q_seq_len, topk] + + pad_heads = heads + if heads < 16: + pad_heads = 16 + + @T.prim_func + def tl_indexer_bwd_kernel( + IndexQ: T.Tensor(index_q_shape, dtype), + IndexK: T.Tensor(index_k_shape, dtype), + Weights: T.Tensor(weights_shape, FP32), + TopkIndices: T.Tensor(topk_indices_shape, INT32), + OGrad: T.Tensor(shape_p, FP32), + dIndexQ: T.Tensor(index_q_shape, dtype), + dWeights: T.Tensor(weights_shape, FP32), + dIndexK: T.Tensor(index_k_shape, FP32), + ): + + with T.Kernel(q_seq_len, threads=num_threads) as (bx): + index_q_shared = T.alloc_shared([pad_heads, dim], dtype=FP32) + weights_shared = T.alloc_shared([pad_heads], dtype=FP32) + index_k_shared = T.alloc_shared([block_I, dim], dtype=FP32) + indices_shared = T.alloc_shared([block_I], dtype=INT32) + d_index_q_frag = T.alloc_fragment([pad_heads, dim], dtype=accum_dtype) + d_weights_frag = T.alloc_fragment([pad_heads], dtype=accum_dtype) + d_index_k_frag = T.alloc_fragment([block_I, dim], dtype=accum_dtype) + logits = T.alloc_fragment((block_I, pad_heads), dtype=accum_dtype) + _logits = T.alloc_shared((block_I, pad_heads), dtype=accum_dtype) + grad = T.alloc_shared([block_I], dtype=FP32) + + num_blocks = T.ceildiv(topk, block_I) + for i, j in T.Parallel(pad_heads, dim): + index_q_shared[i, j] = T.if_then_else(i < heads, IndexQ[bx, i, j], 0) + for i in T.Parallel(heads): + weights_shared[i] = Weights[bx, i] + + T.fill(d_index_q_frag, 0) + T.fill(d_weights_frag, 0) + + for bi_i in T.serial(num_blocks): + for i in T.Parallel(block_I): + if bi_i * block_I + i < topk: + indices_shared[i] = TopkIndices[bx, bi_i * block_I + i] + grad[i] = OGrad[bx, bi_i * block_I + i] + + T.sync_threads() + for i, j in T.Parallel(block_I, dim): + index_k_shared[i, j] = T.if_then_else( + indices_shared[i] > -1 and indices_shared[i] < seq_len, + IndexK[indices_shared[i], j], + 0, + ) + + T.sync_threads() + T.gemm( + index_k_shared, + index_q_shared, + logits, + transpose_A=False, + transpose_B=True, + clear_accum=True, + ) + d_weights_i = T.alloc_fragment((block_I, pad_heads), accum_dtype) + for i, j in T.Parallel(block_I, heads): + d_weights_i[i, j] = grad[i] * ( + T.max(logits[i, j], 0) if use_relu else logits[i, j] + ) + T.reduce_sum(d_weights_i, d_weights_frag, dim=0, clear=False) + + for i, j in T.Parallel(block_I, pad_heads): + _logits[i, j] = T.if_then_else( + (logits[i, j] > 0 if use_relu else True) and j < heads, + grad[i] * weights_shared[j], + 0, + ) + T.sync_threads() + T.gemm( + _logits, + index_k_shared, + d_index_q_frag, + transpose_A=True, + transpose_B=False, + clear_accum=False, + ) + + T.gemm( + _logits, + index_q_shared, + d_index_k_frag, + transpose_A=False, + transpose_B=False, + clear_accum=True, + ) + + for i, j in T.Parallel(block_I, dim): + if indices_shared[i] > -1 and indices_shared[i] < seq_len: + T.atomic_add(dIndexK[indices_shared[i], j], d_index_k_frag[i, j]) + + T.copy(d_index_q_frag[:heads, :], dIndexQ[bx, :, :]) + T.copy(d_weights_frag[:heads], dWeights[bx, :]) + + return tl_indexer_bwd_kernel + + +def indexer_bwd_interface( # pragma: no cover + index_q: torch.Tensor, + weights: torch.Tensor, + index_k: torch.Tensor, + topk_indices: torch.Tensor, + grad_scores: torch.Tensor, + use_relu: bool = True, +): + """Run indexer backward kernel and return gradients for q/w/k.""" + require_tilelang() + _, head_num, head_dim = index_q.shape + k_top = int(topk_indices.shape[1]) + assert k_top > 0, "topk must be positive" + padded_topk = _canonical_topk(k_top) + + if padded_topk != k_top: + padded_indices = torch.full( + (topk_indices.size(0), padded_topk), + -1, + dtype=topk_indices.dtype, + device=topk_indices.device, + ) + padded_indices[:, :k_top].copy_(topk_indices) + topk_indices = padded_indices + + padded_grad_scores = torch.zeros( + (grad_scores.size(0), padded_topk), dtype=grad_scores.dtype, device=grad_scores.device + ) + padded_grad_scores[:, :k_top].copy_(grad_scores) + grad_scores = padded_grad_scores + + grad_scores = grad_scores.contiguous() + weights_kernel = weights.to(dtype=torch.float32).contiguous() + grad_q = torch.empty_like(index_q) + grad_w = torch.empty_like(weights, dtype=torch.float32) + grad_k = torch.zeros_like(index_k, dtype=torch.float32) + + bwd_kernel = _get_indexer_bwd_kernel(head_num, head_dim, padded_topk, use_relu=use_relu) + bwd_kernel( + index_q.contiguous(), + index_k.contiguous(), + weights_kernel, + topk_indices.contiguous(), + grad_scores, + grad_q, + grad_w, + grad_k, + ) + + return grad_q, grad_w, grad_k.to(index_k.dtype) diff --git a/megatron/core/transformer/experimental_attention_variant/ops/tilelang_indexer_fwd.py b/megatron/core/transformer/experimental_attention_variant/ops/tilelang_indexer_fwd.py new file mode 100644 index 00000000000..ffa45ccdf90 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/ops/tilelang_indexer_fwd.py @@ -0,0 +1,224 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# ruff: noqa +# Adapted from: +# https://github.com/tile-ai/tilelang/blob/4956b5835fa554af6c03d4a6289cad44bf310869/ +# examples/deepseek_v32/fp8_lighting_indexer.py +import threading +from collections import OrderedDict + +import torch + +from .tilelang_utils import ( + HAVE_TILELANG, + T, + _get_cached_kernel, + require_tilelang, + tilelang, + tilelang_jit, +) + +_tilelang_indexer_fwd_kernel_cache = OrderedDict() +_tilelang_indexer_clean_logits_kernel_cache = OrderedDict() +_tilelang_indexer_fwd_cache_lock = threading.Lock() + + +def _get_clean_logits_kernel(threads: int = 512, block_K: int = 4096): + return _get_cached_kernel( + _tilelang_indexer_clean_logits_kernel_cache, + _tilelang_indexer_fwd_cache_lock, + (threads, block_K), + lambda: clean_logits_(threads=threads, block_K=block_K), + ) + + +def _get_indexer_fwd_kernel( + heads: int, + index_dim: int, + block_N: int = 256, + num_stages: int = 3, + threads: int = 512, + use_relu: bool = True, +): + return _get_cached_kernel( + _tilelang_indexer_fwd_kernel_cache, + _tilelang_indexer_fwd_cache_lock, + (heads, index_dim, block_N, num_stages, threads, use_relu), + lambda: tl_indexer_fwd_impl( + heads=heads, + index_dim=index_dim, + block_N=block_N, + num_stages=num_stages, + threads=threads, + use_relu=use_relu, + ), + ) + + +_TL_INDEXER_FWD_PASS_CONFIGS = ( + {tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True} if HAVE_TILELANG else {} +) + + +@tilelang_jit(pass_configs=_TL_INDEXER_FWD_PASS_CONFIGS) +def tl_indexer_fwd_impl( # pragma: no cover + heads, index_dim, block_N=256, num_stages=3, threads=512, block_Q=None, use_relu=True +): + """Build tilelang forward kernel for sparse indexer logits.""" + require_tilelang() + assert heads > 0 + if block_Q is None: + block_Q = max(1, 128 // heads) + dtype = T.bfloat16 + accum_dtype = T.float32 + index_dtype = T.int32 + + seq_len = T.dynamic("seq_len") + seq_len_kv = T.dynamic("seq_len_kv") + + index_q_shape = [seq_len * heads, index_dim] + index_k_shape = [seq_len_kv, index_dim] + logits_shape = [seq_len, seq_len_kv] + + @T.prim_func + def tl_indexer_fwd_kernel( + IndexQ: T.Tensor(index_q_shape, dtype), # type: ignore + IndexK: T.Tensor(index_k_shape, dtype), # type: ignore + Logits: T.Tensor(logits_shape, accum_dtype), # type: ignore + Weights: T.Tensor([seq_len, heads], accum_dtype), # type: ignore + CuSeqLenKS: T.Tensor([seq_len], index_dtype), # type: ignore + CuSeqLenKE: T.Tensor([seq_len], index_dtype), # type: ignore + ): + with T.Kernel(T.ceildiv(seq_len, block_Q), threads=threads) as bx: + index_q_shared = T.alloc_shared([block_Q * heads, index_dim], dtype) + index_k_shared = T.alloc_shared([block_N, index_dim], dtype) + s = T.alloc_fragment([block_N, block_Q * heads], accum_dtype) + s_reshaped = T.reshape(s, (block_N, block_Q, heads)) + logits_shared = T.alloc_shared([block_N, block_Q], accum_dtype) + weights = T.alloc_fragment([block_Q, heads], accum_dtype) + + seq_len_i = bx * block_Q + + cu_k_s_min = T.alloc_var(index_dtype) + cu_k_e_max = T.alloc_var(index_dtype) + + cu_k_s_min = 2147483647 + cu_k_e_max = -2147483648 + + for bq_i in T.serial(block_Q): + q_idx = seq_len_i + bq_i + if q_idx < seq_len: + k_s = T.max(T.min(CuSeqLenKS[q_idx], seq_len_kv), 0) + cu_k_s_min = T.min(cu_k_s_min, k_s) + for bq_i in T.serial(block_Q): + q_idx = seq_len_i + bq_i + if q_idx < seq_len: + k_e = T.max(T.min(CuSeqLenKE[q_idx], seq_len_kv), 0) + cu_k_e_max = T.max(cu_k_e_max, k_e) + + # Clamp bounds to [0, seq_len_kv] and normalize empty rows. + cu_k_s_min = T.max(cu_k_s_min, 0) + cu_k_s_min = T.min(cu_k_s_min, seq_len_kv) + cu_k_e_max = T.max(cu_k_e_max, 0) + cu_k_e_max = T.min(cu_k_e_max, seq_len_kv) + if cu_k_e_max < cu_k_s_min: + cu_k_e_max = cu_k_s_min + + for bq_i, h_i, d_i in T.Parallel(block_Q, heads, index_dim): + q_idx = seq_len_i + bq_i + index_q_shared[bq_i * heads + h_i, d_i] = T.if_then_else( + q_idx < seq_len, IndexQ[q_idx * heads + h_i, d_i], 0 + ) + for bq_i, h_i in T.Parallel(block_Q, heads): + q_idx = seq_len_i + bq_i + weights[bq_i, h_i] = T.if_then_else(q_idx < seq_len, Weights[q_idx, h_i], 0) + + for nbn_i in T.Pipelined( + T.ceildiv(cu_k_e_max - cu_k_s_min, block_N), num_stages=num_stages + ): + for bn_i, d_i in T.Parallel(block_N, index_dim): + k_idx = cu_k_s_min + nbn_i * block_N + bn_i + index_k_shared[bn_i, d_i] = T.if_then_else( + k_idx >= 0 and k_idx < cu_k_e_max, IndexK[k_idx, d_i], 0 + ) + + T.gemm( + index_k_shared, + index_q_shared, + s, + transpose_B=True, + clear_accum=True, + policy=T.GemmWarpPolicy.FullCol, + ) + + for bn_i, bq_i, h_i in T.Parallel(block_N, block_Q, heads): + s_reshaped[bn_i, bq_i, h_i] = ( + T.max(s_reshaped[bn_i, bq_i, h_i], 0) + if use_relu + else s_reshaped[bn_i, bq_i, h_i] + ) * weights[bq_i, h_i] + + T.reduce_sum(s_reshaped, logits_shared, dim=-1, clear=True) + + # Keep this write deterministic to satisfy data-race verification. + for bq_i in T.serial(block_Q): + q_idx = seq_len_i + bq_i + if q_idx < seq_len: + for bn_i in T.serial(block_N): + k_idx = cu_k_s_min + nbn_i * block_N + bn_i + if k_idx >= 0 and k_idx < cu_k_e_max: + Logits[q_idx, k_idx] = logits_shared[bn_i, bq_i] + + return tl_indexer_fwd_kernel + + +@tilelang_jit +def clean_logits_(threads: int = 512, block_K: int = 4096): # pragma: no cover + """Build kernel that masks out invalid key ranges in logits.""" + require_tilelang() + seq_len = T.dynamic("seq_len") + seq_len_kv = T.dynamic("seq_len_kv") + + dtype = T.float + indices_dtype = T.int32 + + @T.prim_func + def clean_logits_kernel( + Logits: T.Tensor([seq_len, seq_len_kv], dtype), # type: ignore + CuSeqLenKS: T.Tensor([seq_len], indices_dtype), # type: ignore + CuSeqLenKE: T.Tensor([seq_len], indices_dtype), # type: ignore + ): + with T.Kernel(seq_len, threads=threads) as bx: + tx = T.thread_binding(0, threads, thread="threadIdx.x") + cu_k_s = CuSeqLenKS[bx] + cu_k_e = CuSeqLenKE[bx] + + for n_i in T.Pipelined(T.ceildiv(seq_len_kv, block_K)): + for k_i in T.serial(block_K // threads): + idx = n_i * block_K + k_i * threads + tx + if idx < seq_len_kv and (idx < cu_k_s or idx >= cu_k_e): + Logits[bx, idx] = -T.infinity(dtype) + + return clean_logits_kernel + + +def indexer_fwd_interface( # pragma: no cover + q, kv, weights, cu_seqlen_ks, cu_seqlen_ke, clean_logits=True, use_relu=True +): + """Run indexer forward kernel and optionally clean logits by row bounds.""" + require_tilelang() + seq_len, heads, index_dim = q.shape + seq_len_kv = kv.shape[0] + weights = weights.to(dtype=torch.float32).contiguous() + + tl_indexer_fwd_kernel = _get_indexer_fwd_kernel( + heads=heads, index_dim=index_dim, use_relu=use_relu + ) + logits = torch.empty([seq_len, seq_len_kv], device=q.device, dtype=torch.float32) + tl_indexer_fwd_kernel( + q.view(seq_len * heads, index_dim), kv, logits, weights, cu_seqlen_ks, cu_seqlen_ke + ) + + if clean_logits: + clean_logits_kernel = _get_clean_logits_kernel() + clean_logits_kernel(logits, cu_seqlen_ks, cu_seqlen_ke) + return logits diff --git a/megatron/core/transformer/experimental_attention_variant/ops/tilelang_indexer_loss.py b/megatron/core/transformer/experimental_attention_variant/ops/tilelang_indexer_loss.py new file mode 100644 index 00000000000..3bfbfd70c2a --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/ops/tilelang_indexer_loss.py @@ -0,0 +1,344 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""TileLang kernels for the sparse DSA indexer KL target and score gradient.""" + +import threading +from collections import OrderedDict + +import torch + +from .tilelang_utils import ( + HAVE_TILELANG, + T, + _get_cached_kernel, + _normalize_sm_scale, + require_tilelang, + tilelang, + tilelang_jit, +) + +_target_kernel_cache = OrderedDict() +_kl_kernel_cache = OrderedDict() +_kernel_cache_lock = threading.Lock() + +_PASS_CONFIGS = {tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True} if HAVE_TILELANG else {} + + +def _get_target_kernel( + heads: int, + dim: int, + topk: int, + softmax_scale: float, + block_h: int = 32, + block_i: int = 64, + num_stages: int = 2, + threads: int = 256, +): + scale = _normalize_sm_scale(softmax_scale) + key = (heads, dim, topk, scale, block_h, block_i, num_stages, threads) + return _get_cached_kernel( + _target_kernel_cache, + _kernel_cache_lock, + key, + lambda: sparse_indexer_target( + heads=heads, + dim=dim, + topk=topk, + softmax_scale=scale, + block_h=block_h, + block_i=block_i, + num_stages=num_stages, + threads=threads, + ), + ) + + +def _get_kl_kernel(topk: int, block_i: int = 256, threads: int = 256): + key = (topk, block_i, threads) + return _get_cached_kernel( + _kl_kernel_cache, + _kernel_cache_lock, + key, + lambda: sparse_indexer_kl(topk=topk, block_i=block_i, threads=threads), + ) + + +@tilelang_jit(out_idx=[-1], pass_configs=_PASS_CONFIGS) +def sparse_indexer_target( # pragma: no cover + heads: int, + dim: int, + topk: int, + softmax_scale: float, + block_h: int = 32, + block_i: int = 64, + num_stages: int = 2, + threads: int = 256, +): + """Build a kernel that sums selected-key attention probabilities over local heads.""" + require_tilelang() + assert heads > 0 + assert dim % 16 == 0 + assert topk % block_i == 0 + + seq_len = T.dynamic("seq_len") + seq_len_kv = T.dynamic("seq_len_kv") + dtype = T.bfloat16 + accum_dtype = T.float32 + index_dtype = T.int32 + num_tiles = tilelang.cdiv(topk, block_i) + num_head_tiles = tilelang.cdiv(heads, block_h) + scale_log2 = softmax_scale * 1.4426950408889634 + + @T.prim_func + def main( + Query: T.Tensor([seq_len, heads, dim], dtype), # type: ignore + Key: T.Tensor([seq_len_kv, dim], dtype), # type: ignore + Indices: T.Tensor([seq_len, topk], index_dtype), # type: ignore + Target: T.Tensor([seq_len, topk], accum_dtype), # type: ignore + ): + with T.Kernel(seq_len, threads=threads) as row: + query_shared = T.alloc_shared([block_h, dim], dtype) + key_shared = T.alloc_shared([block_i, dim], dtype) + scores = T.alloc_fragment([block_h, block_i], accum_dtype) + probabilities = T.alloc_fragment([block_h, block_i], accum_dtype) + target_tile = T.alloc_fragment([block_i], accum_dtype) + valid = T.alloc_fragment([block_i], "bool") + row_max = T.alloc_fragment([block_h], accum_dtype) + previous_max = T.alloc_fragment([block_h], accum_dtype) + tile_max = T.alloc_fragment([block_h], accum_dtype) + row_sum = T.alloc_fragment([block_h], accum_dtype) + tile_sum = T.alloc_fragment([block_h], accum_dtype) + alpha = T.alloc_fragment([block_h], accum_dtype) + + for item in T.Parallel(topk): + Target[row, item] = 0 + + for head_tile in T.serial(num_head_tiles): + for head, d in T.Parallel(block_h, dim): + head_index = head_tile * block_h + head + query_shared[head, d] = T.if_then_else( + head_index < heads, Query[row, head_index, d], 0 + ) + T.fill(row_max, -(2**30)) + T.fill(row_sum, 0) + + for tile in T.Pipelined(num_tiles, num_stages=num_stages): + for item in T.Parallel(block_i): + index = Indices[row, tile * block_i + item] + valid[item] = index >= 0 and index < seq_len_kv + for item, d in T.Parallel(block_i, dim): + index = Indices[row, tile * block_i + item] + safe_index = T.max(T.min(index, seq_len_kv - 1), 0) + key_shared[item, d] = T.if_then_else(valid[item], Key[safe_index, d], 0) + T.gemm( + query_shared, + key_shared, + scores, + transpose_B=True, + clear_accum=True, + policy=T.GemmWarpPolicy.FullRow, + ) + for head, item in T.Parallel(block_h, block_i): + scores[head, item] = T.if_then_else( + valid[item] and head_tile * block_h + head < heads, + scores[head, item], + -T.infinity(accum_dtype), + ) + T.copy(row_max, previous_max) + T.reduce_max(scores, tile_max, dim=1, clear=True) + for head in T.Parallel(block_h): + row_max[head] = T.max(previous_max[head], tile_max[head]) + alpha[head] = T.exp2((previous_max[head] - row_max[head]) * scale_log2) + for head, item in T.Parallel(block_h, block_i): + probabilities[head, item] = T.if_then_else( + valid[item] and head_tile * block_h + head < heads, + T.exp2((scores[head, item] - row_max[head]) * scale_log2), + 0, + ) + T.reduce_sum(probabilities, tile_sum, dim=1, clear=True) + for head in T.Parallel(block_h): + row_sum[head] = row_sum[head] * alpha[head] + tile_sum[head] + + for tile in T.Pipelined(num_tiles, num_stages=num_stages): + for item in T.Parallel(block_i): + index = Indices[row, tile * block_i + item] + valid[item] = index >= 0 and index < seq_len_kv + for item, d in T.Parallel(block_i, dim): + index = Indices[row, tile * block_i + item] + safe_index = T.max(T.min(index, seq_len_kv - 1), 0) + key_shared[item, d] = T.if_then_else(valid[item], Key[safe_index, d], 0) + T.gemm( + query_shared, + key_shared, + scores, + transpose_B=True, + clear_accum=True, + policy=T.GemmWarpPolicy.FullRow, + ) + for head, item in T.Parallel(block_h, block_i): + scores[head, item] = T.if_then_else( + valid[item] and head_tile * block_h + head < heads, + scores[head, item], + -T.infinity(accum_dtype), + ) + probabilities[head, item] = T.if_then_else( + valid[item] + and head_tile * block_h + head < heads + and row_sum[head] > 0, + T.exp2((scores[head, item] - row_max[head]) * scale_log2) + / row_sum[head], + 0, + ) + T.reduce_sum(probabilities, target_tile, dim=0, clear=True) + for item in T.Parallel(block_i): + Target[row, tile * block_i + item] += target_tile[item] + + return main + + +@tilelang_jit(out_idx=[-2, -1], pass_configs=_PASS_CONFIGS) +def sparse_indexer_kl(topk: int, block_i: int = 256, threads: int = 256): # pragma: no cover + """Build a kernel that computes sparse KL row sums and gradients for indexer logits.""" + require_tilelang() + assert topk % block_i == 0 + + seq_len = T.dynamic("seq_len") + accum_dtype = T.float32 + num_tiles = tilelang.cdiv(topk, block_i) + log2_e = 1.4426950408889634 + ln_2 = 0.6931471805599453 + eps = 1.0e-10 + + @T.prim_func + def main( + Target: T.Tensor([seq_len, topk], accum_dtype), # type: ignore + IndexLogits: T.Tensor([seq_len, topk], accum_dtype), # type: ignore + ValidMask: T.Tensor([seq_len, topk], "bool"), # type: ignore + GradLogits: T.Tensor([seq_len, topk], accum_dtype), # type: ignore + KLRows: T.Tensor([seq_len], accum_dtype), # type: ignore + ): + with T.Kernel(seq_len, threads=threads) as row: + logits = T.alloc_fragment([1, block_i], accum_dtype) + target = T.alloc_fragment([1, block_i], accum_dtype) + probabilities = T.alloc_fragment([1, block_i], accum_dtype) + kl_terms = T.alloc_fragment([1, block_i], accum_dtype) + valid = T.alloc_fragment([block_i], "bool") + row_max = T.alloc_fragment([1], accum_dtype) + previous_max = T.alloc_fragment([1], accum_dtype) + tile_max = T.alloc_fragment([1], accum_dtype) + row_sum = T.alloc_fragment([1], accum_dtype) + tile_sum = T.alloc_fragment([1], accum_dtype) + target_sum = T.alloc_fragment([1], accum_dtype) + target_tile_sum = T.alloc_fragment([1], accum_dtype) + kl_sum = T.alloc_fragment([1], accum_dtype) + kl_tile_sum = T.alloc_fragment([1], accum_dtype) + + T.fill(row_max, -(2**30)) + T.fill(row_sum, 0) + T.fill(target_sum, 0) + T.fill(kl_sum, 0) + + for tile in T.serial(num_tiles): + for item in T.Parallel(block_i): + valid[item] = ValidMask[row, tile * block_i + item] + logits[0, item] = T.if_then_else( + valid[item], + IndexLogits[row, tile * block_i + item], + -T.infinity(accum_dtype), + ) + target[0, item] = T.if_then_else( + valid[item], Target[row, tile * block_i + item], 0 + ) + T.copy(row_max, previous_max) + T.reduce_max(logits, tile_max, dim=1, clear=True) + row_max[0] = T.max(previous_max[0], tile_max[0]) + for item in T.Parallel(block_i): + probabilities[0, item] = T.if_then_else( + valid[item], T.exp2((logits[0, item] - row_max[0]) * log2_e), 0 + ) + T.reduce_sum(probabilities, tile_sum, dim=1, clear=True) + row_sum[0] = ( + row_sum[0] * T.exp2((previous_max[0] - row_max[0]) * log2_e) + tile_sum[0] + ) + T.reduce_sum(target, target_tile_sum, dim=1, clear=True) + target_sum[0] += target_tile_sum[0] + + for tile in T.serial(num_tiles): + for item in T.Parallel(block_i): + valid[item] = ValidMask[row, tile * block_i + item] + logits[0, item] = T.if_then_else( + valid[item], + IndexLogits[row, tile * block_i + item], + -T.infinity(accum_dtype), + ) + target[0, item] = T.if_then_else( + valid[item] and target_sum[0] > 0, + Target[row, tile * block_i + item] / target_sum[0], + 0, + ) + probabilities[0, item] = T.if_then_else( + valid[item] and row_sum[0] > 0, + T.exp2((logits[0, item] - row_max[0]) * log2_e) / row_sum[0], + 0, + ) + GradLogits[row, tile * block_i + item] = T.if_then_else( + valid[item], probabilities[0, item] - target[0, item], 0 + ) + kl_terms[0, item] = T.if_then_else( + valid[item] and target[0, item] > 0, + target[0, item] + * ( + T.log2(T.max(target[0, item], eps)) * ln_2 + - (logits[0, item] - row_max[0]) + + T.log2(T.max(row_sum[0], eps)) * ln_2 + ), + 0, + ) + T.reduce_sum(kl_terms, kl_tile_sum, dim=1, clear=True) + kl_sum[0] += kl_tile_sum[0] + + KLRows[row] = kl_sum[0] + + return main + + +def sparse_indexer_target_interface( + query: torch.Tensor, key: torch.Tensor, topk_indices: torch.Tensor, softmax_scale: float +) -> torch.Tensor: + """Compute the local-head sparse attention target on selected top-k keys.""" + require_tilelang() + seq_len, heads, dim = query.shape + topk = topk_indices.size(1) + kernel = _get_target_kernel(heads, dim, topk, softmax_scale) + return kernel(query, key, topk_indices) + + +def sparse_indexer_kl_interface( + target: torch.Tensor, index_logits: torch.Tensor, valid_mask: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute unscaled indexer KL sum and its exact gradient with respect to logits.""" + require_tilelang() + kernel = _get_kl_kernel(valid_mask.size(1)) + grad_logits, kl_rows = kernel(target, index_logits, valid_mask) + return kl_rows.sum(), grad_logits + + +class SparseIndexerKLLoss(torch.autograd.Function): # pragma: no cover + """Autograd bridge from fused sparse KL score gradients to the TileLang indexer.""" + + @staticmethod + def forward(ctx, target, index_logits, valid_mask): + """Compute the sparse indexer KL loss and save its logits gradient.""" + kl_sum, grad_logits = sparse_indexer_kl_interface(target, index_logits, valid_mask) + ctx.save_for_backward(grad_logits) + return kl_sum + + @staticmethod + def backward(ctx, grad_output): + """Scale the saved index-logits gradient for the backward pass.""" + (grad_logits,) = ctx.saved_tensors + return None, grad_logits * grad_output, None + + +if not HAVE_TILELANG: + SparseIndexerKLLoss = None diff --git a/megatron/core/transformer/experimental_attention_variant/ops/tilelang_sparse_mla_bwd.py b/megatron/core/transformer/experimental_attention_variant/ops/tilelang_sparse_mla_bwd.py new file mode 100644 index 00000000000..1cccec8339b --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/ops/tilelang_sparse_mla_bwd.py @@ -0,0 +1,529 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# ruff: noqa +# Adapted from: +# https://github.com/tile-ai/tilelang/blob/4ff81c7d40803d269569e157e847623e84553f78/ +# examples/deepseek_v32/sparse_mla_bwd.py +import threading +from collections import OrderedDict + +import torch + +from .tilelang_utils import ( + HAVE_TILELANG, + T, + _env_int, + _get_cached_kernel, + _normalize_sm_scale, + _round_up, + require_tilelang, + tilelang, + tilelang_jit, +) + +_SPARSE_MLA_BWD_BLOCK_SIZE = 32 +_tilelang_sparse_mla_preprocess_kernel_cache = OrderedDict() +_tilelang_sparse_mla_bwd_kernel_cache = OrderedDict() +_tilelang_sparse_mla_postprocess_kernel_cache = OrderedDict() +_tilelang_sparse_mla_bwd_cache_lock = threading.Lock() + + +def _get_preprocess_kernel(H: int, D: int): + return _get_cached_kernel( + _tilelang_sparse_mla_preprocess_kernel_cache, + _tilelang_sparse_mla_bwd_cache_lock, + (H, D), + lambda: preprocess(H, D), + ) + + +def _normalize_block_h(block_h: int) -> int: + if block_h >= 64: + return 64 + if block_h >= 32: + return 32 + return 16 + + +def _get_bwd_kernel( + H: int, D: int, D_tail: int, topk: int, kv_group: int, sm_scale, max_block_h: int +): + max_block_h = _normalize_block_h(max_block_h) + key = (H, D, D_tail, topk, kv_group, _normalize_sm_scale(sm_scale), max_block_h) + return _get_cached_kernel( + _tilelang_sparse_mla_bwd_kernel_cache, + _tilelang_sparse_mla_bwd_cache_lock, + key, + lambda: bwd(H, D, D_tail, topk, kv_group, sm_scale, max_block_h=max_block_h), + ) + + +def _get_postprocess_kernel(D: int, D_tail: int, kv_group: int): + return _get_cached_kernel( + _tilelang_sparse_mla_postprocess_kernel_cache, + _tilelang_sparse_mla_bwd_cache_lock, + (D, D_tail, kv_group), + lambda: postprocess(D, D_tail, kv_group), + ) + + +@tilelang_jit(out_idx=[-1]) +def preprocess( # pragma: no cover + H, + D, + block_ND=32, + num_stages=5, + dtype=T.bfloat16 if HAVE_TILELANG else None, + accum_dtype=T.float32 if HAVE_TILELANG else None, +): + """Build preprocessing kernel that computes Delta = sum(O * dO) per row/head.""" + require_tilelang() + assert dtype == T.bfloat16 + assert accum_dtype == T.float32 + batch = T.dynamic("batch") + seq_len = T.dynamic("seq_len") + shape = [batch, seq_len, H, D] + + @T.prim_func + def preprocess_kernel( + O: T.Tensor(shape, dtype), + dO: T.Tensor(shape, dtype), + Delta: T.Tensor([batch, seq_len, H], accum_dtype), + ): + with T.Kernel(H, T.ceildiv(seq_len, block_ND), batch) as (bx, by, bz): + o = T.alloc_fragment([block_ND, block_ND], accum_dtype) + do = T.alloc_fragment([block_ND, block_ND], accum_dtype) + delta = T.alloc_fragment([block_ND], accum_dtype) + acc = T.alloc_fragment([block_ND, block_ND], accum_dtype) + T.clear(acc) + for k in T.Pipelined(T.ceildiv(D, block_ND), num_stages=num_stages): + T.copy( + O[ + bz, + by * block_ND : (by + 1) * block_ND, + bx, + k * block_ND : (k + 1) * block_ND, + ], + o, + ) + T.copy( + dO[ + bz, + by * block_ND : (by + 1) * block_ND, + bx, + k * block_ND : (k + 1) * block_ND, + ], + do, + ) + for i, j in T.Parallel(block_ND, block_ND): + acc[i, j] += o[i, j] * do[i, j] + T.reduce_sum(acc, delta, 1) + T.copy(delta, Delta[bz, by * block_ND : (by + 1) * block_ND, bx]) + + return preprocess_kernel + + +@tilelang_jit(out_idx=[-1]) +def postprocess( # pragma: no cover + D, + D_tail, + kv_group=1, + block_N=64, + threads=128, + dtype=T.bfloat16 if HAVE_TILELANG else None, + accum_dtype=T.float32 if HAVE_TILELANG else None, +): + """Build postprocess kernel that casts/exports accumulated dKV.""" + require_tilelang() + assert dtype == T.bfloat16 + assert accum_dtype == T.float32 + batch = T.dynamic("batch") + seq_len_kv = T.dynamic("seq_len_kv") + dkv_shape = [batch, seq_len_kv, kv_group, D + D_tail] + + @T.prim_func + def postprocess_kernel( + dKV: T.Tensor(dkv_shape, accum_dtype), dKV_out: T.Tensor(dkv_shape, dtype) + ): + with T.Kernel(T.ceildiv(seq_len_kv, block_N), kv_group, batch, threads=threads) as ( + bx, + by, + bz, + ): + T.copy( + dKV[bz, bx * block_N : (bx + 1) * block_N, by, :], + dKV_out[bz, bx * block_N : (bx + 1) * block_N, by, :], + ) + + return postprocess_kernel + + +_SPARSE_MLA_BWD_PASS_CONFIGS = ( + { + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + tilelang.PassConfigKey.TL_ENABLE_AGGRESSIVE_SHARED_MEMORY_MERGE: True, + } + if HAVE_TILELANG + else {} +) + + +@tilelang_jit(out_idx=[-2], pass_configs=_SPARSE_MLA_BWD_PASS_CONFIGS) +def bwd( # pragma: no cover + H, + D, + D_tail, + topk, + kv_group=1, + sm_scale=None, + block_size=32, + max_block_h=32, + num_stages=2, + threads=128, + indices_dtype=T.int32 if HAVE_TILELANG else None, + dtype=T.bfloat16 if HAVE_TILELANG else None, + accum_dtype=T.float32 if HAVE_TILELANG else None, +): + """Build sparse-MLA backward kernel.""" + require_tilelang() + assert ( + topk % block_size == 0 + ), "otherwise will load some index=0 thus causing wrong kv to be loaded" + assert dtype == T.bfloat16 + assert accum_dtype == T.float32 + assert indices_dtype == T.int32 + + if sm_scale is None: + sm_scale = (D + D_tail) ** (-0.5) + sm_scale_mul_reciprocal_log2 = sm_scale * 1.44269504 # log2(e) + + batch = T.dynamic("batch") + seq_len = T.dynamic("seq_len") + seq_len_kv = T.dynamic("seq_len_kv") + + H_kv = H // kv_group + q_shape = [batch, seq_len, H, D + D_tail] + k_shape = [batch, seq_len_kv, kv_group, D + D_tail] + o_shape = [batch, seq_len, H, D] + indices_shape = [batch, seq_len, kv_group, topk] + delta_shape = [batch, seq_len, H] + lse_shape = [batch, seq_len, H] + assert indices_dtype == T.int32 + assert dtype == T.bfloat16 + assert accum_dtype == T.float32 + + H = H_kv + padded_H = max(tilelang.math.next_power_of_2(H_kv), 16) + block_H = min(_normalize_block_h(max_block_h), padded_H) + assert padded_H % block_H == 0 + NH = padded_H // block_H + BS = block_size + NS = tilelang.cdiv(topk, block_size) + + split_store = 2 + + @T.prim_func + def sparse_mla_bwd_kernel( + Q: T.Tensor(q_shape, dtype), + KV: T.Tensor(k_shape, dtype), + dO: T.Tensor(o_shape, dtype), + Indices: T.Tensor(indices_shape, indices_dtype), + Lse: T.Tensor(lse_shape, accum_dtype), + Delta: T.Tensor(delta_shape, accum_dtype), + dQ: T.Tensor(q_shape, dtype), + dKV: T.Tensor(k_shape, accum_dtype), + ): + with T.Kernel(seq_len, batch, kv_group * NH, threads=threads) as (s_i, by, bz): + Q_shared = T.alloc_shared([block_H, D], dtype) + Q_tail_shared = T.alloc_shared([block_H, D_tail], dtype) + KV_shared = T.alloc_shared([BS, D], dtype) + KV_tail_shared = T.alloc_shared([BS, D_tail], dtype) + dO_shared = T.alloc_shared([block_H, D], dtype) + mask = T.alloc_fragment([BS], "bool") + + P_shared_cast = T.alloc_shared([block_H, BS], dtype) + dP_shared_cast = T.alloc_shared([block_H, BS], dtype) + dQ_shared = T.alloc_shared([block_H, D], dtype) + dQ_tail_shared = T.alloc_shared([block_H, D_tail], dtype) + + acc_p = T.alloc_fragment([block_H, BS], accum_dtype) + acc_dp = T.alloc_fragment([block_H, BS], accum_dtype) + acc_dq = T.alloc_fragment([block_H, D], accum_dtype) + acc_dq_tail = T.alloc_fragment([block_H, D_tail], accum_dtype) + acc_dkv = T.alloc_fragment([BS, D], accum_dtype) + acc_dkv_tail = T.alloc_fragment([BS, D_tail], accum_dtype) + acc_dkv_shared = T.alloc_shared([BS // split_store, D], accum_dtype) + acc_dkv_tail_shared = T.alloc_shared([BS // split_store, D_tail], accum_dtype) + + T.copy(Q[by, s_i, bz * block_H : (bz + 1) * block_H, :D], Q_shared) + T.copy(Q[by, s_i, bz * block_H : (bz + 1) * block_H, D:], Q_tail_shared) + T.copy(dO[by, s_i, bz * block_H : (bz + 1) * block_H, :D], dO_shared) + + T.clear(acc_dq) + T.clear(acc_dq_tail) + + # Process each block of indices + for i_i in T.Pipelined(NS, num_stages=num_stages): + # Check which indices are valid + for bi_i in T.Parallel(BS): + # Changed here for thd + mask[bi_i] = Indices[by, s_i, bz // NH, i_i * BS + bi_i] != -1 + + # Compute attention scores + for h_i, bi_i in T.Parallel(block_H, BS): + acc_p[h_i, bi_i] = T.if_then_else(mask[bi_i], 0, -T.infinity(acc_p.dtype)) + + # Load KV, V for this block of indices + for bi_i, d_i in T.Parallel(BS, D): + idx = Indices[by, s_i, bz // NH, i_i * BS + bi_i] + safe_idx = T.max(idx, 0) + KV_shared[bi_i, d_i] = KV[by, safe_idx, bz // NH, d_i] + + T.gemm( + Q_shared, KV_shared, acc_p, transpose_B=True, policy=T.GemmWarpPolicy.FullCol + ) + + for bi_i, d_i in T.Parallel(BS, D_tail): + idx = Indices[by, s_i, bz // NH, i_i * BS + bi_i] + safe_idx = T.max(idx, 0) + KV_tail_shared[bi_i, d_i] = KV[by, safe_idx, bz // NH, D + d_i] + T.gemm( + Q_tail_shared, + KV_tail_shared, + acc_p, + transpose_B=True, + policy=T.GemmWarpPolicy.FullCol, + ) + + for h_i, bi_i in T.Parallel(block_H, BS): + acc_p[h_i, bi_i] = T.exp2( + acc_p[h_i, bi_i] * sm_scale_mul_reciprocal_log2 + - Lse[by, s_i, bz * block_H + h_i] + ) + + T.copy(acc_p, P_shared_cast) + + T.gemm( + dO_shared, + KV_shared, + acc_dp, + transpose_B=True, + policy=T.GemmWarpPolicy.FullCol, + clear_accum=True, + ) + + for h_i, bi_i in T.Parallel(block_H, BS): + acc_dp[h_i, bi_i] = ( + acc_p[h_i, bi_i] + * (acc_dp[h_i, bi_i] - Delta[by, s_i, bz * block_H + h_i]) + * sm_scale + ) + + T.copy(acc_dp, dP_shared_cast) + T.gemm(dP_shared_cast, KV_shared, acc_dq, policy=T.GemmWarpPolicy.FullCol) + T.gemm(dP_shared_cast, KV_tail_shared, acc_dq_tail, policy=T.GemmWarpPolicy.FullCol) + + T.gemm( + dP_shared_cast, + Q_shared, + acc_dkv, + transpose_A=True, + policy=T.GemmWarpPolicy.FullCol, + clear_accum=True, + ) + T.gemm( + P_shared_cast, + dO_shared, + acc_dkv, + transpose_A=True, + policy=T.GemmWarpPolicy.FullCol, + ) + + T.clear(acc_dkv_tail) + T.gemm( + dP_shared_cast, + Q_tail_shared, + acc_dkv_tail, + transpose_A=True, + policy=T.GemmWarpPolicy.FullCol, + ) + + for s in range(split_store): + for bi_i, d_i in T.Parallel(BS, D): + if bi_i < BS // split_store: + acc_dkv_shared[bi_i, d_i] = acc_dkv[bi_i + s * (BS // split_store), d_i] + + for bi_i, d_i in T.Parallel(BS, D_tail): + if bi_i < BS // split_store: + acc_dkv_tail_shared[bi_i, d_i] = acc_dkv_tail[ + bi_i + s * (BS // split_store), d_i + ] + + for bi_i, d_i in T.Parallel(BS // split_store, D // 4): + idx = Indices[by, s_i, bz // NH, i_i * BS + bi_i + s * (BS // split_store)] + if idx >= 0: + T.atomic_addx4( + dKV[by, idx, bz // NH, d_i * 4], acc_dkv_shared[bi_i, d_i * 4] + ) + + # Atomically update dKV, dKV_tail tensors + for bi_i, d_i in T.Parallel(BS // split_store, D_tail // 4): + idx = Indices[by, s_i, bz // NH, i_i * BS + bi_i + s * (BS // split_store)] + if idx >= 0: + T.atomic_addx4( + dKV[by, idx, bz // NH, D + d_i * 4], + acc_dkv_tail_shared[bi_i, d_i * 4], + ) + + # Store the accumulated dQ + T.copy(acc_dq, dQ_shared) + T.copy(acc_dq_tail, dQ_tail_shared) + + T.copy(dQ_shared, dQ[by, s_i, bz * block_H : (bz + 1) * block_H, :D]) + T.copy(dQ_tail_shared, dQ[by, s_i, bz * block_H : (bz + 1) * block_H, D:]) + + return sparse_mla_bwd_kernel + + +def _sparse_mla_delta_batched(o, do): # pragma: no cover + """Compute Delta = sum(O * dO) with safe sequence padding for TileLang tiles.""" + require_tilelang() + assert o.is_contiguous() + assert do.is_contiguous() + assert o.shape == do.shape + B, S, H, D = o.shape + + seq_len_padded = _round_up(S, _SPARSE_MLA_BWD_BLOCK_SIZE) + if seq_len_padded != S: + o_padded = torch.zeros((B, seq_len_padded, H, D), dtype=o.dtype, device=o.device) + o_padded[:, :S].copy_(o) + o = o_padded + + do_padded = torch.zeros((B, seq_len_padded, H, D), dtype=do.dtype, device=do.device) + do_padded[:, :S].copy_(do) + do = do_padded + + preprocess_kernel = _get_preprocess_kernel(H, D) + return preprocess_kernel(o, do)[:, :S].contiguous() + + +def sparse_mla_delta(o, do): # pragma: no cover + """Compute Delta = sum(O * dO) per sequence row and head.""" + squeeze_batch = o.ndim == 3 + if squeeze_batch: + o = o.unsqueeze(0) + do = do.unsqueeze(0) + delta = _sparse_mla_delta_batched(o, do) + if squeeze_batch: + delta = delta.squeeze(0) + return delta + + +def sparse_mla_bwd(q, kv, o, do, indices, lse, sm_scale=None, delta=None): # pragma: no cover + """Run sparse-MLA backward kernels and return (dq, dkv).""" + require_tilelang() + + seq_bucket = _env_int("MCORE_DSA_TILELANG_SEQ_BUCKET", 256) + topk_bucket = _env_int("MCORE_DSA_TILELANG_TOPK_BUCKET", _SPARSE_MLA_BWD_BLOCK_SIZE) + max_block_h = _env_int("MCORE_DSA_TILELANG_BWD_MAX_BLOCK_H", 32) + + squeeze_batch = q.ndim == 3 + if squeeze_batch: + q = q.unsqueeze(0) + kv = kv.unsqueeze(0) + do = do.unsqueeze(0) + indices = indices.unsqueeze(0) + lse = lse.unsqueeze(0) + if o is not None: + if squeeze_batch: + o = o.unsqueeze(0) + + assert q.is_contiguous() + assert kv.is_contiguous() + assert indices.is_contiguous() + assert lse.is_contiguous() + assert q.ndim == 4 and kv.ndim == 4 and do.ndim == 4 and indices.ndim == 4 and lse.ndim == 3 + B, S, H, dim_plus_tail_dim = q.shape + _, S_kv, kv_group, _ = kv.shape + assert kv.shape[-1] == dim_plus_tail_dim + assert kv.shape[0] == B + # This copied kernel currently assumes a fixed base value-channel dimension. + D = 512 + assert ( + dim_plus_tail_dim >= D + ), f"Invalid dimensions: dim_plus_tail_dim={dim_plus_tail_dim} is smaller than base D={D}" + + D_tail = dim_plus_tail_dim - D + topk = indices.shape[-1] + assert indices.shape == (B, S, kv_group, topk) + assert lse.shape == (B, S, H) + + seq_bucketed = _round_up(S, seq_bucket) + seq_kv_bucketed = _round_up(S_kv, seq_bucket) + topk_bucketed = _round_up(_round_up(topk, topk_bucket), _SPARSE_MLA_BWD_BLOCK_SIZE) + + if seq_bucketed != S: + q_padded = torch.zeros( + (B, seq_bucketed, H, dim_plus_tail_dim), dtype=q.dtype, device=q.device + ) + q_padded[:, :S].copy_(q) + q = q_padded + + if o is not None: + o_padded = torch.zeros((B, seq_bucketed, H, D), dtype=o.dtype, device=o.device) + o_padded[:, :S].copy_(o) + o = o_padded + + do_padded = torch.zeros((B, seq_bucketed, H, D), dtype=do.dtype, device=do.device) + do_padded[:, :S].copy_(do) + do = do_padded + + lse_padded = torch.zeros((B, seq_bucketed, H), dtype=lse.dtype, device=lse.device) + lse_padded[:, :S].copy_(lse) + lse = lse_padded + + if seq_kv_bucketed != S_kv: + kv_padded = torch.zeros( + (B, seq_kv_bucketed, kv_group, dim_plus_tail_dim), dtype=kv.dtype, device=kv.device + ) + kv_padded[:, :S_kv].copy_(kv) + kv = kv_padded + + if seq_bucketed != S or topk_bucketed != topk: + indices_padded = torch.full( + (B, seq_bucketed, kv_group, topk_bucketed), + -1, + dtype=indices.dtype, + device=indices.device, + ) + indices_padded[:, :S, :, :topk].copy_(indices) + indices = indices_padded + + if delta is not None: + if delta.ndim == 2: + delta = delta.unsqueeze(0) + if seq_bucketed != S: + delta_padded = torch.zeros((B, seq_bucketed, H), dtype=delta.dtype, device=delta.device) + delta_padded[:, :S].copy_(delta) + delta = delta_padded + + # Get kernels + bwd_kernel = _get_bwd_kernel(H, D, D_tail, topk_bucketed, kv_group, sm_scale, max_block_h) + postprocess_kernel = _get_postprocess_kernel(D, D_tail, kv_group) + + if delta is None: + if o is None: + raise ValueError("sparse_mla_bwd requires either output tensor o or precomputed delta") + delta = _sparse_mla_delta_batched(o, do) + dkv = torch.zeros_like(kv, dtype=torch.float32) + dq = bwd_kernel(q, kv, do, indices, lse, delta, dkv) + dkv = postprocess_kernel(dkv) + + dq = dq[:, :S].contiguous() + dkv = dkv[:, :S_kv].contiguous() + + if squeeze_batch: + dq = dq.squeeze(0) + dkv = dkv.squeeze(0) + + return dq, dkv diff --git a/megatron/core/transformer/experimental_attention_variant/ops/tilelang_sparse_mla_fwd.py b/megatron/core/transformer/experimental_attention_variant/ops/tilelang_sparse_mla_fwd.py new file mode 100644 index 00000000000..707c453d00e --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/ops/tilelang_sparse_mla_fwd.py @@ -0,0 +1,310 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# ruff: noqa +# Adapted from: +# https://github.com/tile-ai/tilelang/blob/e666d2d3cc483829c57618c9ebf2e4f4ada0819d/ +# examples/deepseek_v32/sparse_mla_fwd.py +import threading +from collections import OrderedDict + +import torch + +from .tilelang_utils import ( + HAVE_TILELANG, + T, + _env_int, + _get_cached_kernel, + _normalize_sm_scale, + _round_up, + require_tilelang, + tilelang, + tilelang_jit, +) + +_tilelang_sparse_mla_fwd_kernel_cache = OrderedDict() +_tilelang_sparse_mla_fwd_cache_lock = threading.Lock() + + +def _get_sparse_mla_fwd_kernel( + heads: int, + dim: int, + tail_dim: int, + topk: int, + kv_group: int, + sm_scale, + block_I: int, + num_stages: int, + threads: int, +): + key = ( + heads, + dim, + tail_dim, + topk, + kv_group, + _normalize_sm_scale(sm_scale), + block_I, + num_stages, + threads, + ) + return _get_cached_kernel( + _tilelang_sparse_mla_fwd_kernel_cache, + _tilelang_sparse_mla_fwd_cache_lock, + key, + lambda: sparse_mla_fwd( + heads, + dim, + tail_dim, + topk, + kv_group, + sm_scale, + block_I=block_I, + num_stages=num_stages, + threads=threads, + ), + ) + + +_SPARSE_MLA_FWD_PASS_CONFIGS = ( + { + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + } + if HAVE_TILELANG + else {} +) + + +@tilelang_jit(out_idx=[-2, -1], pass_configs=_SPARSE_MLA_FWD_PASS_CONFIGS) +def sparse_mla_fwd( # pragma: no cover + heads, dim, tail_dim, topk, kv_group=1, sm_scale=None, block_I=64, num_stages=2, threads=256 +): + """Build sparse-MLA forward kernel.""" + require_tilelang() + assert dim == tilelang.math.next_power_of_2(dim), f"dim must be a power of two, got dim={dim}" + assert tail_dim == tilelang.math.next_power_of_2( + tail_dim + ), f"tail_dim must be a power of two, got tail_dim={tail_dim}" + assert ( + topk % block_I == 0 + ), "otherwise will load some index=0 thus causing wrong kv to be loaded" + if sm_scale is None: + sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e) + else: + sm_scale = sm_scale * 1.44269504 # log2(e) + + batch = T.dynamic("batch") + seq_len = T.dynamic("seq_len") + seq_len_kv = T.dynamic("seq_len_kv") + + head_kv = heads // kv_group + q_shape = [batch, seq_len, heads, dim + tail_dim] + kv_shape = [batch, seq_len_kv, kv_group, dim + tail_dim] + o_shape = [batch, seq_len, heads, dim] + indices_shape = [batch, seq_len, kv_group, topk] + lse_shape = [batch, seq_len, heads] + indices_dtype = T.int32 + dtype = T.bfloat16 + accum_dtype = T.float32 + + G = kv_group + H = head_kv + padded_H = max(tilelang.math.next_power_of_2(head_kv), 16) + if padded_H != H: + assert kv_group == 1, ( + "here we solve the H padding automatically, otherwise handle Q/Output copy with " + "your own mask (for kv_group==1, g_i*padded_H:(g_i+1)*padded_H is handled)" + ) + BI = block_I + NI = tilelang.cdiv(topk, block_I) + D = dim + D_tail = tail_dim + + if head_kv > 64: + assert head_kv % 64 == 0, "head_kv should be a multiple of 64" + REPLICATE_H = head_kv // 64 + else: + REPLICATE_H = 1 + + H_per_block = padded_H if REPLICATE_H == 1 else 64 + + @T.prim_func + def main( + Q: T.Tensor(q_shape, dtype), # type: ignore + KV: T.Tensor(kv_shape, dtype), # type: ignore + Indices: T.Tensor(indices_shape, indices_dtype), # type: ignore + Output: T.Tensor(o_shape, dtype), # type: ignore + Lse: T.Tensor(lse_shape, accum_dtype), # type: ignore + ): + with T.Kernel(seq_len * REPLICATE_H, batch, kv_group, threads=threads) as (bx, by, bz): + Q_shared = T.alloc_shared([H_per_block, D], dtype) + Q_tail_shared = T.alloc_shared([H_per_block, D_tail], dtype) + KV_shared = T.alloc_shared([BI, D], dtype) + K_tail_shared = T.alloc_shared([BI, D_tail], dtype) + O_shared = T.alloc_shared([H_per_block, D], dtype) + Lse_shared = T.alloc_shared([H_per_block], accum_dtype) + mask = T.alloc_fragment([BI], "bool") + + acc_o = T.alloc_fragment([H_per_block, D], accum_dtype) + acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) + S_shared = T.alloc_shared([H_per_block, BI], dtype) + sumexp = T.alloc_fragment([H_per_block], accum_dtype) + sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) + alpha = T.alloc_fragment([H_per_block], accum_dtype) + m_i = T.alloc_fragment([H_per_block], accum_dtype) + m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) + + T.fill(acc_o, 0) + T.fill(sumexp, 0) + T.fill(m_i, -(2**30)) # avoid -inf - inf to cause nan + + b_i, g_i = by, bz + s_i = bx if REPLICATE_H == 1 else (bx // REPLICATE_H) + q_i = s_i + max_kv_i = q_i + + H0 = g_i * padded_H + (0 if REPLICATE_H == 1 else (bx % REPLICATE_H) * 64) + H1 = H0 + H_per_block + + T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared) + T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared) + + for i_i in T.Pipelined(NI, num_stages=num_stages): + for bi_i in T.Parallel(BI): + # Changed here for thd + mask[bi_i] = Indices[b_i, s_i, g_i, i_i * BI + bi_i] != -1 + + for bi_i, d_i in T.Parallel(BI, D): + idx = Indices[b_i, s_i, g_i, i_i * BI + bi_i] + safe_idx = T.max(idx, 0) + KV_shared[bi_i, d_i] = KV[b_i, safe_idx, g_i, d_i] + for bi_i, d_i in T.Parallel(BI, D_tail): + idx = Indices[b_i, s_i, g_i, i_i * BI + bi_i] + safe_idx = T.max(idx, 0) + K_tail_shared[bi_i, d_i] = KV[b_i, safe_idx, g_i, D + d_i] + + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else(mask[bi_i], 0, -T.infinity(acc_s.dtype)) + T.gemm( + Q_shared, KV_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow + ) + T.gemm( + Q_tail_shared, + K_tail_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.copy(m_i, m_i_prev) + T.reduce_max(acc_s, m_i, dim=1, clear=False) + for h_i in T.Parallel(H_per_block): + m_i[h_i] = T.max(m_i[h_i], m_i_prev[h_i]) + for h_i in T.Parallel(H_per_block): + alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.exp2(acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale) + # Reduce the current tile; the online softmax accumulation happens below. + T.reduce_sum(acc_s, sumexp_i, dim=1) + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] = acc_o[h_i, d_i] * alpha[h_i] + + T.copy(acc_s, S_shared) + T.gemm(S_shared, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) + + # Rescale. Packed THD can produce sentinel-only rows; define those rows as zero + # output/LSE instead of dividing by a zero softmax denominator. + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] = T.if_then_else(sumexp[h_i] > 0, acc_o[h_i, d_i] / sumexp[h_i], 0) + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = T.if_then_else( + sumexp[h_i] > 0, T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale, 0 + ) + + T.copy(acc_o, Output[b_i, s_i, H0:H1, :]) + T.copy(sumexp, Lse[b_i, s_i, H0:H1]) + + return main + + +def sparse_mla_fwd_interface( + q, kv, indices, sm_scale=None, d_v=512, block_I=64, num_stages=2, threads=256 +): + """Run sparse-MLA forward kernel and return (out, lse).""" + require_tilelang() + seq_bucket = _env_int("MCORE_DSA_TILELANG_SEQ_BUCKET", 256) + topk_bucket = _env_int("MCORE_DSA_TILELANG_TOPK_BUCKET", block_I) + + squeeze_batch = q.ndim == 3 + if squeeze_batch: + q = q.unsqueeze(0) + kv = kv.unsqueeze(0) + indices = indices.unsqueeze(0) + + assert q.is_contiguous() and kv.is_contiguous() and indices.is_contiguous() + assert q.ndim == 4 and kv.ndim == 4 and indices.ndim == 4 + batch, seq_len, heads, dim_plus_tail_dim = q.shape + _, seq_len_kv, kv_group, kv_dim = kv.shape + assert ( + kv_dim == dim_plus_tail_dim + ), "q and kv must have the same embedding dimension on the last axis" + assert ( + dim_plus_tail_dim == 576 + ), "TileLang sparse MLA fwd is currently specialized for dim_plus_tail_dim=576" + dim = d_v + assert 0 < dim <= dim_plus_tail_dim, f"d_v must be in (0, {dim_plus_tail_dim}], but got {dim}" + + assert kv.shape[-1] == dim_plus_tail_dim + tail_dim = dim_plus_tail_dim - dim + assert kv.shape[0] == batch + _, _, _, topk = indices.shape + assert indices.shape == (batch, seq_len, kv_group, topk) + + seq_len_bucketed = _round_up(seq_len, seq_bucket) + seq_len_kv_bucketed = _round_up(seq_len_kv, seq_bucket) + topk_bucketed = _round_up(_round_up(topk, topk_bucket), block_I) + + if seq_len_bucketed != seq_len: + q_padded = torch.zeros( + (batch, seq_len_bucketed, heads, dim_plus_tail_dim), dtype=q.dtype, device=q.device + ) + q_padded[:, :seq_len].copy_(q) + q = q_padded + + if seq_len_kv_bucketed != seq_len_kv: + kv_padded = torch.zeros( + (batch, seq_len_kv_bucketed, kv_group, dim_plus_tail_dim), + dtype=kv.dtype, + device=kv.device, + ) + kv_padded[:, :seq_len_kv].copy_(kv) + kv = kv_padded + + if seq_len_bucketed != seq_len or topk_bucketed != topk: + indices_padded = torch.full( + (batch, seq_len_bucketed, kv_group, topk_bucketed), + -1, + dtype=indices.dtype, + device=indices.device, + ) + indices_padded[:, :seq_len, :, :topk].copy_(indices) + indices = indices_padded + + kernel = _get_sparse_mla_fwd_kernel( + heads=heads, + dim=dim, + tail_dim=tail_dim, + topk=topk_bucketed, + kv_group=kv_group, + sm_scale=sm_scale, + block_I=block_I, + num_stages=num_stages, + threads=threads, + ) + out, lse = kernel(q, kv, indices) + out = out[:, :seq_len].contiguous() + lse = lse[:, :seq_len].contiguous() + if squeeze_batch: + out = out.squeeze(0) + lse = lse.squeeze(0) + return out, lse diff --git a/megatron/core/transformer/experimental_attention_variant/ops/tilelang_utils.py b/megatron/core/transformer/experimental_attention_variant/ops/tilelang_utils.py new file mode 100644 index 00000000000..689436bfb59 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/ops/tilelang_utils.py @@ -0,0 +1,101 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import os +from collections import OrderedDict + +import torch + +from megatron.core.utils import round_up_to_nearest_multiple + +try: + import tilelang + from tilelang import language as T # pylint: disable=unused-import + + HAVE_TILELANG = True +except (ImportError, OSError): + tilelang = None + T = None + HAVE_TILELANG = False + + +def _noop_jit(*args, **kwargs): + if len(args) == 1 and callable(args[0]) and not kwargs: + return args[0] + + def decorator(func): + return func + + return decorator + + +def tilelang_jit(*args, **kwargs): + """Return TileLang's jit decorator when available, otherwise a no-op decorator.""" + if HAVE_TILELANG: + return tilelang.jit(*args, **kwargs) + return _noop_jit(*args, **kwargs) + + +def require_tilelang(): + """Raise a clear error when a fused TileLang kernel is used without TileLang installed.""" + if not HAVE_TILELANG: + raise ImportError( + "TileLang is required to use fused DSA TileLang kernels. " + "Install tilelang or use the unfused fallback path." + ) + + +def _env_int(name: str, default: int) -> int: + """Parse a positive integer environment variable, falling back to ``default``.""" + value = os.getenv(name) + if value is None: + return default + try: + parsed = int(value) + except ValueError: + return default + return parsed if parsed > 0 else default + + +_TILELANG_KERNEL_CACHE_MAX = _env_int("MCORE_DSA_TILELANG_KERNEL_CACHE_MAX", 512) + + +def _cache_put_lru(cache: OrderedDict, key, value): + """Insert ``value`` as the most-recently-used entry, evicting oldest past the cap.""" + cache[key] = value + cache.move_to_end(key) + while len(cache) > _TILELANG_KERNEL_CACHE_MAX: + cache.popitem(last=False) + + +def _get_cached_kernel(cache: OrderedDict, lock, key, build_fn): + """Return a cached compiled kernel for ``key``, building it via ``build_fn`` on miss.""" + with lock: + kernel = cache.pop(key, None) + if kernel is None: + kernel = build_fn() + _cache_put_lru(cache, key, kernel) + return kernel + + +def _round_up(x: int, multiple: int) -> int: + if multiple <= 1: + return x + return round_up_to_nearest_multiple(x, multiple) + + +def _next_power_of_two(x: int) -> int: + if x <= 1: + return 1 + return 1 << (x - 1).bit_length() + + +def _normalize_sm_scale(sm_scale): + """Coerce a softmax scale to a stable float so it can key the kernel cache.""" + if sm_scale is None: + return None + if isinstance(sm_scale, torch.Tensor): + sm_scale = float(sm_scale.detach().item()) + else: + sm_scale = float(sm_scale) + # Avoid tiny floating-point jitter creating cache-key churn. + return round(sm_scale, 12) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_tilelang_kernels.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_tilelang_kernels.py new file mode 100644 index 00000000000..2a9e151d4ee --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_tilelang_kernels.py @@ -0,0 +1,1284 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import math +from types import SimpleNamespace + +import pytest +import torch + +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.transformer.experimental_attention_variant import ( + dsa_indexer_loss, + dsa_masking, + dsa_tilelang_kernels, +) +from megatron.core.transformer.experimental_attention_variant.ops import ( + indexer, + sparse_mla, + tilelang_dsa, + tilelang_indexer_bwd, + tilelang_indexer_fwd, + tilelang_indexer_loss, + tilelang_sparse_mla_bwd, + tilelang_utils, +) + + +def test_run_fused_qk_topk_forwards_to_tilelang_backend(monkeypatch): + q = torch.empty(2, 1, 3, 4) + k = torch.empty(5, 1, 4) + weights = torch.empty(2, 1, 3) + starts = torch.tensor([0, 1], dtype=torch.int32) + ends = torch.tensor([3, 5], dtype=torch.int32) + expected_indices = torch.tensor([[[2, 1], [4, 3]]], dtype=torch.int32) + call = {} + + def fake_run_fused_qk_topk( + q_arg, k_arg, weights_arg, index_topk, starts_arg, ends_arg, block_size, use_relu, **kwargs + ): + call.update( + q=q_arg, + k=k_arg, + weights=weights_arg, + index_topk=index_topk, + starts=starts_arg, + ends=ends_arg, + block_size=block_size, + use_relu=use_relu, + kwargs=kwargs, + ) + return expected_indices + + monkeypatch.setattr( + dsa_tilelang_kernels.tilelang_dsa, "run_fused_qk_topk", fake_run_fused_qk_topk + ) + + result = dsa_tilelang_kernels.run_fused_qk_topk( + q, + k, + weights, + index_topk=2, + starts=starts, + ends=ends, + block_size=8, + use_relu=False, + use_local_indexer_varlen=True, + ) + + indices, topk_length = result + assert indices is expected_indices + assert topk_length is None + assert call["q"] is q + assert call["k"] is k + assert call["weights"] is weights + assert call["index_topk"] == 2 + assert call["starts"] is starts + assert call["ends"] is ends + assert call["block_size"] == 8 + assert call["use_relu"] is False + assert call["kwargs"]["use_local_indexer_varlen"] is True + assert call["kwargs"]["cp_size"] == 1 + + +def test_run_fused_qk_topk_preserves_unavailable_backend(monkeypatch): + def fake_run_fused_qk_topk(*_args, **_kwargs): + return None + + monkeypatch.setattr( + dsa_tilelang_kernels.tilelang_dsa, "run_fused_qk_topk", fake_run_fused_qk_topk + ) + + result = dsa_tilelang_kernels.run_fused_qk_topk( + torch.empty(2, 1, 3, 4), + torch.empty(5, 1, 4), + torch.empty(2, 1, 3), + index_topk=2, + starts=torch.tensor([0, 1], dtype=torch.int32), + ends=torch.tensor([3, 5], dtype=torch.int32), + block_size=8, + ) + + assert result is None + + +def test_tilelang_packed_cp_indexer_inputs_segment_keys_and_bounds(): + packed_seq_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=torch.tensor([0, 8, 24], dtype=torch.int32), + cu_seqlens_kv=torch.tensor([0, 8, 24], dtype=torch.int32), + max_seqlen_q=16, + max_seqlen_kv=16, + ) + query_positions = torch.tensor([0, 1, 6, 7, 8, 9, 10, 11, 20, 21, 22, 23]) + starts = torch.tensor([0] * 4 + [8] * 8, dtype=torch.int32) + ends = (query_positions + 1).to(torch.int32) + index_k = torch.arange(24, dtype=torch.float32).view(24, 1) + + segmented_k, local_starts, local_ends, source_indices = ( + tilelang_dsa._build_packed_cp_indexer_inputs( + index_k, + starts, + ends, + packed_seq_params=packed_seq_params, + cp_size=2, + cp_rank=0, + single_packed_thd_sequence=False, + local_query_start=0, + local_query_len=12, + ) + ) + + expected_sources = torch.tensor( + [0, 1, *range(8), *range(8, 12), *range(8, 24)], dtype=torch.int64 + ) + torch.testing.assert_close(source_indices, expected_sources) + torch.testing.assert_close(segmented_k[:, 0], expected_sources.to(torch.float32)) + torch.testing.assert_close( + local_starts, torch.tensor([0, 0, 2, 2, 10, 10, 10, 10, 14, 14, 14, 14], dtype=torch.int32) + ) + torch.testing.assert_close( + local_ends, torch.tensor([1, 2, 9, 10, 11, 12, 13, 14, 27, 28, 29, 30], dtype=torch.int32) + ) + + +def test_tilelang_packed_cp_indexer_remaps_segmented_topk(monkeypatch): + packed_seq_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=torch.tensor([0, 8, 24], dtype=torch.int32), + cu_seqlens_kv=torch.tensor([0, 8, 24], dtype=torch.int32), + max_seqlen_q=16, + max_seqlen_kv=16, + ) + query_positions = torch.tensor([0, 1, 6, 7, 8, 9, 10, 11, 20, 21, 22, 23]) + starts = torch.tensor([0] * 4 + [8] * 8, dtype=torch.int32) + ends = (query_positions + 1).to(torch.int32) + seen = {} + + def fake_lighting_indexer_indices( + index_q, index_k, index_w, starts_arg, ends_arg, index_topk, use_relu=True + ): + del index_q, index_w, use_relu + seen["key"] = index_k[:, 0].clone() + seen["starts"] = starts_arg.clone() + seen["ends"] = ends_arg.clone() + offsets = torch.arange(index_topk, dtype=torch.int32).view(1, -1) + return ends_arg.view(-1, 1) - 1 - offsets + + monkeypatch.setattr(tilelang_dsa, "lighting_indexer_indices", fake_lighting_indexer_indices) + topk = tilelang_dsa.fused_qk_topk_lighting( + torch.ones((12, 1, 1, 1), dtype=torch.bfloat16), + torch.arange(24, dtype=torch.bfloat16).view(24, 1, 1), + torch.ones((12, 1, 1)), + index_topk=3, + starts=starts, + ends=ends, + block_size=12, + use_local_indexer_varlen=True, + packed_seq_params=packed_seq_params, + cp_size=2, + ) + + expected = [] + for position, sequence_start in zip(query_positions.tolist(), starts.tolist()): + row = list(range(position, max(sequence_start - 1, position - 3), -1)) + expected.append(row + [-1] * (3 - len(row))) + torch.testing.assert_close(topk, torch.tensor([expected], dtype=torch.int32)) + assert seen["key"].numel() == 30 + torch.testing.assert_close( + seen["starts"], + torch.tensor([0, 0, 2, 2, 10, 10, 10, 10, 14, 14, 14, 14], dtype=torch.int32), + ) + + +def test_run_fused_qk_topk_with_loss_preserves_unavailable_backend(monkeypatch): + def fake_run_fused_qk_topk_with_loss(**kwargs): + return None + + monkeypatch.setattr( + dsa_tilelang_kernels.tilelang_dsa, + "run_fused_qk_topk_with_loss", + fake_run_fused_qk_topk_with_loss, + ) + + result = dsa_tilelang_kernels.run_fused_qk_topk_with_loss( + q=torch.empty(2, 1, 3, 4), + k=torch.empty(5, 1, 4), + weights=torch.empty(2, 1, 3), + index_topk=2, + starts=torch.tensor([0, 1], dtype=torch.int32), + ends=torch.tensor([3, 5], dtype=torch.int32), + block_size=8, + query=torch.empty(2, 1, 3, 4), + key=torch.empty(5, 1, 1, 4), + softmax_scale=0.5, + loss_coeff=0.1, + pg_collection=SimpleNamespace(), + config=SimpleNamespace(), + use_local_indexer_varlen=True, + ) + + assert result is None + + +def test_run_fused_qk_topk_with_loss_adds_empty_topk_length(monkeypatch): + q = torch.empty(2, 1, 3, 4) + k = torch.empty(5, 1, 4) + weights = torch.empty(2, 1, 3) + starts = torch.tensor([0, 1], dtype=torch.int32) + ends = torch.tensor([3, 5], dtype=torch.int32) + query = torch.empty(2, 1, 3, 4) + key = torch.empty(5, 1, 1, 4) + query_valid_rows = torch.tensor([[True, False]]) + pg_collection = SimpleNamespace() + expected_indices = torch.tensor([[[2, 1], [4, 3]]], dtype=torch.int32) + expected_loss = torch.tensor(1.25) + call = {} + + def fake_run_fused_qk_topk_with_loss(**kwargs): + call.update(kwargs) + return expected_indices, expected_loss + + monkeypatch.setattr( + dsa_tilelang_kernels.tilelang_dsa, + "run_fused_qk_topk_with_loss", + fake_run_fused_qk_topk_with_loss, + ) + + result = dsa_tilelang_kernels.run_fused_qk_topk_with_loss( + q=q, + k=k, + weights=weights, + index_topk=2, + starts=starts, + ends=ends, + block_size=8, + query=query, + key=key, + softmax_scale=0.5, + loss_coeff=0.1, + pg_collection=pg_collection, + query_valid_rows=query_valid_rows, + calculate_per_token_loss=True, + use_relu=False, + config=SimpleNamespace(), + use_local_indexer_varlen=True, + ) + + indices, topk_length, indexer_loss = result + assert indices is expected_indices + assert topk_length is None + assert indexer_loss is expected_loss + assert call["q"] is q + assert call["k"] is k + assert call["weights"] is weights + assert call["index_topk"] == 2 + assert call["starts"] is starts + assert call["ends"] is ends + assert call["block_size"] == 8 + assert call["query"] is query + assert call["key"] is key + assert call["softmax_scale"] == 0.5 + assert call["loss_coeff"] == 0.1 + assert call["pg_collection"] is pg_collection + assert call["query_valid_rows"] is query_valid_rows + assert call["calculate_per_token_loss"] is True + assert call["use_relu"] is False + + +def test_run_fused_absorbed_sparse_attention_forwards_to_tilelang_backend(monkeypatch): + query = torch.empty(2, 1, 3, 4) + key = torch.empty(5, 1, 1, 4) + topk_indices = torch.tensor([[[0, 1], [1, 99]]], dtype=torch.int32) + topk_length = torch.tensor([[2, 1]], dtype=torch.int32) + expected_output = torch.empty(2, 1, 3, 4) + call = {} + + def fake_run_fused_absorbed_sparse_attention( + query_arg, key_arg, topk_indices_arg, softmax_scale, v_channels + ): + call.update( + query=query_arg, + key=key_arg, + topk_indices=topk_indices_arg, + softmax_scale=softmax_scale, + v_channels=v_channels, + ) + return expected_output + + monkeypatch.setattr( + dsa_tilelang_kernels.tilelang_dsa, + "run_fused_absorbed_sparse_attention", + fake_run_fused_absorbed_sparse_attention, + ) + + result = dsa_tilelang_kernels.run_fused_absorbed_sparse_attention( + query, key, topk_indices, softmax_scale=0.5, v_channels=4, topk_length=topk_length + ) + + assert result is expected_output + assert call["query"] is query + assert call["key"] is key + torch.testing.assert_close( + call["topk_indices"], torch.tensor([[[0, 1], [1, -1]]], dtype=torch.int32) + ) + assert call["softmax_scale"] == 0.5 + assert call["v_channels"] == 4 + + +def test_indexer_topk_helpers_mask_invalid_entries(): + logits = torch.tensor([[1.0, 3.0, float("-inf")], [0.0, 2.0, 1.0]]) + requested_indices = torch.tensor([[1, -1, 4], [0, 2, 1]], dtype=torch.int32) + + gathered = indexer.pytorch_extract_topk_scores(logits, requested_indices) + + assert torch.equal(gathered[0], torch.tensor([3.0, float("-inf"), float("-inf")])) + assert torch.equal(gathered[1], torch.tensor([0.0, 1.0, 2.0])) + + topk_scores, topk_indices = indexer._select_topk_from_logits(logits, topk=4) + assert topk_scores.shape == (2, 3) + assert topk_indices.shape == (2, 3) + assert topk_indices.dtype == torch.int32 + assert -1 in topk_indices[0].tolist() + + empty_scores, empty_indices = indexer._select_topk_from_logits(torch.empty(2, 0), topk=3) + assert empty_scores.shape == (2, 0) + assert empty_indices.shape == (2, 0) + assert empty_indices.dtype == torch.int32 + + +def test_sparse_mla_head_mask_helpers(): + indices = torch.tensor([[[0, -1], [-1, -1]], [[1, 2], [3, -1]]], dtype=torch.int32) + + valid_heads = sparse_mla._valid_head_mask(indices, num_heads=4) + + assert torch.equal( + valid_heads, torch.tensor([[True, True, False, False], [True, True, True, True]]) + ) + + tensor = torch.arange(2 * 4 * 3, dtype=torch.float32).view(2, 4, 3) + zeroed = sparse_mla._zero_invalid_heads(tensor, valid_heads) + + assert torch.equal(zeroed[0, :2], tensor[0, :2]) + assert torch.equal(zeroed[0, 2:], torch.zeros_like(tensor[0, 2:])) + assert torch.equal(zeroed[1], tensor[1]) + + batched_indices = indices.unsqueeze(0) + batched_valid_heads = sparse_mla._valid_head_mask(batched_indices, num_heads=4) + assert torch.equal(batched_valid_heads, valid_heads.unsqueeze(0)) + + batched_tensor = tensor.unsqueeze(0) + batched_zeroed = sparse_mla._zero_invalid_heads(batched_tensor, batched_valid_heads) + assert torch.equal(batched_zeroed, zeroed.unsqueeze(0)) + + with pytest.raises(RuntimeError, match="heads must be divisible"): + sparse_mla._valid_head_mask(indices, num_heads=3) + + +def test_tilelang_dsa_sanitize_helper(): + topk_indices = torch.tensor([[0, 2, 5], [-1, 3, 4]], dtype=torch.int32) + topk_scores = torch.tensor([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]) + starts = torch.tensor([1, 3], dtype=torch.int32) + ends = torch.tensor([5, 4], dtype=torch.int32) + + sanitized_indices, sanitized_scores = tilelang_dsa._sanitize_fused_topk_outputs( + topk_indices, starts, ends, topk_scores + ) + + assert torch.equal(sanitized_indices, torch.tensor([[-1, 2, -1], [-1, 3, -1]])) + assert torch.equal( + torch.isneginf(sanitized_scores), torch.tensor([[True, False, True], [True, False, True]]) + ) + + +def test_tilelang_dsa_scratch_cache_reuses_buffers(monkeypatch): + tilelang_dsa._DSA_SCRATCH_CACHE.clear() + monkeypatch.setattr(tilelang_dsa, "_DSA_SCRATCH_CACHE_TOTAL_BYTES", 0) + monkeypatch.setattr(tilelang_dsa, "_DSA_SCRATCH_CACHE_MAX_ENTRIES", 1) + monkeypatch.setattr(tilelang_dsa, "_DSA_SCRATCH_CACHE_MAX_BYTES", 1024) + + first = tilelang_dsa._get_scratch_buffer("a", (2,), torch.float32, torch.device("cpu")) + first.fill_(3.0) + reused = tilelang_dsa._get_scratch_buffer("a", (2,), torch.float32, torch.device("cpu")) + second = tilelang_dsa._get_scratch_buffer("b", (2,), torch.float32, torch.device("cpu")) + + assert reused is first + assert torch.equal(reused, torch.full((2,), 3.0)) + assert list(tilelang_dsa._DSA_SCRATCH_CACHE) == [ + ("b", (2,), torch.float32, torch.device("cpu")) + ] + assert tilelang_dsa._DSA_SCRATCH_CACHE_TOTAL_BYTES == second.numel() * second.element_size() + + +def test_tilelang_kernel_helper_caches_and_env_parsing(monkeypatch): + monkeypatch.delenv("MCORE_DSA_TILELANG_KERNEL_CACHE_MAX", raising=False) + assert tilelang_utils._env_int("MCORE_DSA_TILELANG_KERNEL_CACHE_MAX", 7) == 7 + + monkeypatch.setenv("MCORE_DSA_TILELANG_KERNEL_CACHE_MAX", "bad") + assert tilelang_utils._env_int("MCORE_DSA_TILELANG_KERNEL_CACHE_MAX", 7) == 7 + + monkeypatch.setenv("MCORE_DSA_TILELANG_KERNEL_CACHE_MAX", "-3") + assert tilelang_utils._env_int("MCORE_DSA_TILELANG_KERNEL_CACHE_MAX", 7) == 7 + + monkeypatch.setenv("MCORE_DSA_TILELANG_KERNEL_CACHE_MAX", "2") + assert tilelang_utils._env_int("MCORE_DSA_TILELANG_KERNEL_CACHE_MAX", 7) == 2 + + # Shared numeric/layout helpers now live in tilelang_utils. + assert tilelang_utils._next_power_of_two(0) == 1 + assert tilelang_utils._next_power_of_two(9) == 16 + assert tilelang_utils._round_up(9, 4) == 12 + assert tilelang_utils._round_up(9, 1) == 9 + assert tilelang_utils._normalize_sm_scale(None) is None + assert tilelang_utils._normalize_sm_scale(torch.tensor(0.5)) == 0.5 + + # Kernel-specific helpers stay with their modules. + assert tilelang_indexer_bwd._canonical_topk(33) == 64 + assert tilelang_indexer_bwd.is_supported_indexer_bwd_head_count(8) + assert tilelang_indexer_bwd.is_supported_indexer_bwd_head_count(64) + assert not tilelang_indexer_bwd.is_supported_indexer_bwd_head_count(7) + assert not tilelang_indexer_bwd.is_supported_indexer_bwd_head_count(72) + assert tilelang_sparse_mla_bwd._normalize_block_h(12) == 16 + assert tilelang_sparse_mla_bwd._normalize_block_h(40) == 32 + assert tilelang_sparse_mla_bwd._normalize_block_h(80) == 64 + assert tilelang_dsa._is_supported_sparse_mla_head_count(16) + assert tilelang_dsa._is_supported_sparse_mla_head_count(32) + assert tilelang_dsa._is_supported_sparse_mla_head_count(64) + assert tilelang_dsa._is_supported_sparse_mla_head_count(128) + assert tilelang_dsa._is_supported_sparse_mla_head_count(256, kv_group=2) + assert not tilelang_dsa._is_supported_sparse_mla_head_count(96) + assert not tilelang_dsa._is_supported_sparse_mla_head_count(96, kv_group=0) + # head_kv that is not a power of two >= 16 pads to a larger head dim in the kernels and + # would index past the real head count, so it must decline to the unfused path. + assert not tilelang_dsa._is_supported_sparse_mla_head_count(8) + assert not tilelang_dsa._is_supported_sparse_mla_head_count(48) + assert not tilelang_dsa._is_supported_sparse_mla_head_count(192) + assert not tilelang_dsa._is_supported_sparse_mla_head_count(96, kv_group=2) + + +def test_sparse_mla_canonicalizes_size_one_batch_stride_without_copy(): + tensor_sbhd = torch.empty(256, 1, 1, 4) + tensor_bshd = tensor_sbhd.permute(1, 0, 2, 3) + + assert tensor_bshd.is_contiguous() + assert tensor_bshd.stride(0) != tensor_bshd.numel() + + normalized = sparse_mla._canonicalize_batch_stride(tensor_bshd) + + assert normalized.stride(0) == normalized.numel() + assert normalized.data_ptr() == tensor_bshd.data_ptr() + + +def test_indexer_bwd_returns_grad_k_in_index_k_dtype(monkeypatch): + captured = {} + + def fake_kernel(index_q, index_k, weights, topk_indices, grad_scores, grad_q, grad_w, grad_k): + del index_q, index_k, weights, topk_indices, grad_scores + captured["grad_k_kernel_dtype"] = grad_k.dtype + grad_q.fill_(1) + grad_w.fill_(2) + grad_k.fill_(3) + + monkeypatch.setattr(tilelang_indexer_bwd, "require_tilelang", lambda: None) + monkeypatch.setattr( + tilelang_indexer_bwd, "_get_indexer_bwd_kernel", lambda *_args, **_kwargs: fake_kernel + ) + + index_q = torch.empty((2, 8, 4), dtype=torch.bfloat16) + index_k = torch.empty((3, 4), dtype=torch.bfloat16) + weights = torch.empty((2, 8), dtype=torch.float32) + topk_indices = torch.zeros((2, 1), dtype=torch.int32) + grad_scores = torch.empty((2, 1), dtype=torch.float32) + + _, _, grad_k = tilelang_indexer_bwd.indexer_bwd_interface( + index_q, weights, index_k, topk_indices, grad_scores + ) + + assert captured["grad_k_kernel_dtype"] == torch.float32 + assert grad_k.dtype == index_k.dtype + torch.testing.assert_close(grad_k.float(), torch.full_like(grad_k, 3, dtype=torch.float32)) + + +def test_sparse_mla_delta_pads_partial_sequence_tile(monkeypatch): + seq_len = 65 + padded_seq_len = 96 + heads = 2 + dim = 4 + o = torch.arange(seq_len * heads * dim, dtype=torch.float32).view(seq_len, heads, dim) + do = torch.full_like(o, 2.0) + + monkeypatch.setattr(tilelang_sparse_mla_bwd, "require_tilelang", lambda: None) + + def fake_get_preprocess_kernel(H, D): + assert H == heads + assert D == dim + + def fake_preprocess_kernel(o_arg, do_arg): + assert o_arg.shape == (1, padded_seq_len, heads, dim) + assert do_arg.shape == (1, padded_seq_len, heads, dim) + assert torch.equal(o_arg[:, :seq_len], o.unsqueeze(0)) + assert torch.equal(do_arg[:, :seq_len], do.unsqueeze(0)) + assert torch.equal(o_arg[:, seq_len:], torch.zeros_like(o_arg[:, seq_len:])) + assert torch.equal(do_arg[:, seq_len:], torch.zeros_like(do_arg[:, seq_len:])) + return torch.arange(padded_seq_len * heads, dtype=torch.float32).view( + 1, padded_seq_len, heads + ) + + return fake_preprocess_kernel + + monkeypatch.setattr( + tilelang_sparse_mla_bwd, "_get_preprocess_kernel", fake_get_preprocess_kernel + ) + + delta = tilelang_sparse_mla_bwd.sparse_mla_delta(o.contiguous(), do.contiguous()) + + assert delta.shape == (seq_len, heads) + assert delta.is_contiguous() + expected = torch.arange(padded_seq_len * heads, dtype=torch.float32).view( + padded_seq_len, heads + )[:seq_len] + torch.testing.assert_close(delta, expected) + + +def test_lighting_indexer_indices_preserves_single_head_weight_axis(monkeypatch): + seen = {} + + def fake_indexer_fwd_interface( + index_q, index_k, weights, cu_seqlen_ks, cu_seqlen_ke, clean_logits, use_relu + ): + del index_q, index_k, cu_seqlen_ks, cu_seqlen_ke + seen["weights_shape"] = weights.shape + seen["clean_logits"] = clean_logits + seen["use_relu"] = use_relu + return torch.arange(6, dtype=torch.float32).view(2, 3) + + monkeypatch.setattr(indexer, "indexer_fwd_interface", fake_indexer_fwd_interface) + + topk_indices = indexer.lighting_indexer_indices( + index_q=torch.empty(2, 1, 4), + index_k=torch.empty(3, 4), + weights=torch.ones(2, 1), + cu_seqlen_ks=torch.zeros(2, dtype=torch.int32), + cu_seqlen_ke=torch.full((2,), 3, dtype=torch.int32), + topk=2, + use_relu=False, + ) + + assert seen["weights_shape"] == (2, 1) + assert seen["clean_logits"] is True + assert seen["use_relu"] is False + torch.testing.assert_close(topk_indices, torch.tensor([[2, 1], [2, 1]], dtype=torch.int32)) + + +def _skip_if_real_tilelang_indexer_unavailable(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for TileLang indexer parity tests") + if not indexer.HAVE_TILELANG_INDEXER: + pytest.skip("TileLang indexer forward/backward kernels are unavailable") + + +def _pytorch_indexer_scores(index_q, index_k, weights, *, use_relu): + per_head_scores = torch.einsum("qhd,kd->qkh", index_q.float(), index_k.float()) + if use_relu: + per_head_scores = per_head_scores.relu() + return (per_head_scores * weights.float().unsqueeze(1)).sum(dim=-1) + + +@pytest.mark.parametrize("use_relu", [False, True]) +def test_tilelang_indexer_forward_matches_pytorch(use_relu): + _skip_if_real_tilelang_indexer_unavailable() + torch.manual_seed(1234) + + device = torch.device("cuda") + q_len, k_len, heads, dim = 5, 19, 8, 16 + index_q = (torch.randn(q_len, heads, dim, device=device) * 0.25).to(torch.bfloat16) + index_k = (torch.randn(k_len, dim, device=device) * 0.25).to(torch.bfloat16) + weights = torch.randn(q_len, heads, dtype=torch.float32, device=device) * 0.25 + starts = torch.tensor([0, 1, 3, 5, 8], dtype=torch.int32, device=device) + ends = torch.tensor([7, 10, 13, 17, 19], dtype=torch.int32, device=device) + + actual = indexer.indexer_fwd_interface( + index_q, index_k, weights, starts, ends, clean_logits=True, use_relu=use_relu + ) + expected = _pytorch_indexer_scores(index_q, index_k, weights, use_relu=use_relu) + key_positions = torch.arange(k_len, device=device) + valid = (key_positions.unsqueeze(0) >= starts.unsqueeze(1)) & ( + key_positions.unsqueeze(0) < ends.unsqueeze(1) + ) + + torch.testing.assert_close(actual[valid], expected[valid], rtol=2e-2, atol=2e-2) + assert torch.isneginf(actual[~valid]).all() + + +@pytest.mark.parametrize("use_relu", [False, True]) +def test_tilelang_indexer_backward_matches_pytorch(use_relu): + _skip_if_real_tilelang_indexer_unavailable() + torch.manual_seed(5678) + + device = torch.device("cuda") + q_len, k_len, heads, dim = 4, 32, 8, 16 + index_q = (torch.randn(q_len, heads, dim, device=device) * 0.25).to(torch.bfloat16) + index_k = (torch.randn(k_len, dim, device=device) * 0.25).to(torch.bfloat16) + weights = torch.randn(q_len, heads, dtype=torch.float32, device=device) * 0.25 + topk_indices = torch.tensor( + [ + [0, 2, 4, 6, 8, 10, -1], + [1, 3, 5, 7, 9, 11, -1], + [12, 14, 16, 18, 20, 22, -1], + [13, 15, 17, 19, 21, 23, -1], + ], + dtype=torch.int32, + device=device, + ) + grad_scores = torch.randn(topk_indices.shape, dtype=torch.float32, device=device) + grad_scores.masked_fill_(topk_indices < 0, 0.0) + + actual_grad_q, actual_grad_w, actual_grad_k = indexer.indexer_bwd_interface( + index_q, weights, index_k, topk_indices, grad_scores, use_relu=use_relu + ) + + reference_q = index_q.detach().clone().requires_grad_(True) + reference_k = index_k.detach().clone().requires_grad_(True) + reference_w = weights.detach().clone().requires_grad_(True) + reference_scores = _pytorch_indexer_scores( + reference_q, reference_k, reference_w, use_relu=use_relu + ) + valid = topk_indices >= 0 + selected_scores = reference_scores.gather(1, topk_indices.clamp_min(0).long()) + selected_scores = selected_scores.masked_fill(~valid, 0.0) + (selected_scores * grad_scores).sum().backward() + + torch.testing.assert_close(actual_grad_q, reference_q.grad, rtol=5e-2, atol=5e-2) + torch.testing.assert_close(actual_grad_w, reference_w.grad, rtol=5e-2, atol=5e-2) + torch.testing.assert_close(actual_grad_k, reference_k.grad, rtol=5e-2, atol=5e-2) + + +def test_shared_topk_sort_uses_explicit_validity_mask(): + indices = torch.tensor([[5, 1, 7, 3]], dtype=torch.int32) + scores = torch.tensor([[0.5, 0.1, 0.7, 0.3]]) + valid = torch.tensor([[True, False, True, False]]) + + sorted_indices, sorted_scores = dsa_masking.sort_topk_by_index( + indices, valid, sk=8, topk_scores=scores + ) + + torch.testing.assert_close(sorted_indices, torch.tensor([[5, 7, -1, -1]], dtype=torch.int32)) + torch.testing.assert_close(sorted_scores[:, :2], torch.tensor([[0.5, 0.7]])) + assert torch.isneginf(sorted_scores[:, 2:]).all() + + +def test_tilelang_kernel_getters_reuse_cached_builders(monkeypatch): + def make_fake_kernel(): + return lambda *_args, **_kwargs: None + + try: + monkeypatch.setattr(tilelang_utils, "_TILELANG_KERNEL_CACHE_MAX", 1) + tilelang_indexer_fwd._tilelang_indexer_fwd_kernel_cache.clear() + tilelang_indexer_fwd._tilelang_indexer_clean_logits_kernel_cache.clear() + fwd_builds = [] + clean_builds = [] + + def fake_indexer_builder(**kwargs): + fwd_builds.append(kwargs) + return make_fake_kernel() + + def fake_clean_builder(**kwargs): + clean_builds.append(kwargs) + return make_fake_kernel() + + monkeypatch.setattr(tilelang_indexer_fwd, "tl_indexer_fwd_impl", fake_indexer_builder) + monkeypatch.setattr(tilelang_indexer_fwd, "clean_logits_", fake_clean_builder) + + first = tilelang_indexer_fwd._get_indexer_fwd_kernel(2, 4) + second = tilelang_indexer_fwd._get_indexer_fwd_kernel(2, 4) + third = tilelang_indexer_fwd._get_indexer_fwd_kernel(4, 4) + clean_first = tilelang_indexer_fwd._get_clean_logits_kernel() + clean_second = tilelang_indexer_fwd._get_clean_logits_kernel() + + assert first is second + assert third is not first + assert len(fwd_builds) == 2 + assert clean_first is clean_second + assert len(clean_builds) == 1 + + tilelang_indexer_bwd._tilelang_indexer_bwd_kernel_cache.clear() + bwd_builds = [] + + def fake_bwd_builder(*args, **kwargs): + bwd_builds.append((args, kwargs)) + return make_fake_kernel() + + monkeypatch.setattr(tilelang_indexer_bwd, "tl_indexer_bwd_impl", fake_bwd_builder) + bwd_first = tilelang_indexer_bwd._get_indexer_bwd_kernel(8, 4, 32) + bwd_second = tilelang_indexer_bwd._get_indexer_bwd_kernel(8, 4, 32) + bwd_third = tilelang_indexer_bwd._get_indexer_bwd_kernel(16, 4, 32) + + assert bwd_first is bwd_second + assert bwd_third is not bwd_first + assert len(bwd_builds) == 2 + assert bwd_builds[0][1]["num_threads"] == 32 + assert bwd_builds[1][1]["num_threads"] == 128 + + tilelang_sparse_mla_bwd._tilelang_sparse_mla_preprocess_kernel_cache.clear() + tilelang_sparse_mla_bwd._tilelang_sparse_mla_bwd_kernel_cache.clear() + tilelang_sparse_mla_bwd._tilelang_sparse_mla_postprocess_kernel_cache.clear() + monkeypatch.setattr( + tilelang_sparse_mla_bwd, "preprocess", lambda *args, **kwargs: make_fake_kernel() + ) + monkeypatch.setattr( + tilelang_sparse_mla_bwd, "bwd", lambda *args, **kwargs: make_fake_kernel() + ) + monkeypatch.setattr( + tilelang_sparse_mla_bwd, "postprocess", lambda *args, **kwargs: make_fake_kernel() + ) + + preprocess_first = tilelang_sparse_mla_bwd._get_preprocess_kernel(2, 4) + preprocess_second = tilelang_sparse_mla_bwd._get_preprocess_kernel(2, 4) + sparse_bwd_first = tilelang_sparse_mla_bwd._get_bwd_kernel(2, 512, 64, 32, 1, 0.5, 80) + sparse_bwd_second = tilelang_sparse_mla_bwd._get_bwd_kernel(2, 512, 64, 32, 1, 0.5, 80) + postprocess_first = tilelang_sparse_mla_bwd._get_postprocess_kernel(512, 64, 1) + postprocess_second = tilelang_sparse_mla_bwd._get_postprocess_kernel(512, 64, 1) + + assert preprocess_first is preprocess_second + assert sparse_bwd_first is sparse_bwd_second + assert postprocess_first is postprocess_second + finally: + tilelang_indexer_fwd._tilelang_indexer_fwd_kernel_cache.clear() + tilelang_indexer_fwd._tilelang_indexer_clean_logits_kernel_cache.clear() + tilelang_indexer_bwd._tilelang_indexer_bwd_kernel_cache.clear() + tilelang_sparse_mla_bwd._tilelang_sparse_mla_preprocess_kernel_cache.clear() + tilelang_sparse_mla_bwd._tilelang_sparse_mla_bwd_kernel_cache.clear() + tilelang_sparse_mla_bwd._tilelang_sparse_mla_postprocess_kernel_cache.clear() + + +def test_tilelang_utils_noop_jit_and_require_tilelang(monkeypatch): + def fn(): + return "ok" + + monkeypatch.setattr(tilelang_utils, "HAVE_TILELANG", False) + assert tilelang_utils._noop_jit(fn) is fn + assert tilelang_utils._noop_jit()(fn) is fn + assert tilelang_utils.tilelang_jit(fn) is fn + with pytest.raises(ImportError, match="TileLang is required"): + tilelang_utils.require_tilelang() + + +def test_compute_topk_target_chunk_sum_shared_and_per_head_paths(monkeypatch): + tilelang_dsa._DSA_SCRATCH_CACHE.clear() + monkeypatch.setattr(tilelang_dsa, "_DSA_SCRATCH_CACHE_TOTAL_BYTES", 0) + + query_h = torch.tensor( + [[[1.0, 0.0], [0.0, 1.0]], [[1.0, 1.0], [1.0, -1.0]]], requires_grad=True + ) + key_shared = torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]], requires_grad=True) + idx_seq = torch.tensor([[0, 1], [1, 2]], dtype=torch.int64) + valid_seq = torch.tensor([[True, True], [True, False]]) + + shared = tilelang_dsa._compute_topk_target_chunk_sum( + query_h=query_h, + key_shared=key_shared, + key_per_head=None, + s0=0, + s1=2, + idx_seq=idx_seq, + valid_seq=valid_seq, + softmax_scale=1.0, + head_chunk_size=1, + topk_chunk_size=1, + sk=3, + hn=2, + ) + + assert shared.shape == (2, 2) + assert not shared.requires_grad + assert torch.allclose(shared.sum(dim=-1), torch.tensor([2.0, 2.0]), atol=1e-6) + assert shared[1, 1] == 0 + + key_per_head = torch.stack((key_shared, key_shared + 1.0)).detach().requires_grad_(True) + per_head = tilelang_dsa._compute_topk_target_chunk_sum( + query_h=query_h, + key_shared=None, + key_per_head=key_per_head, + s0=0, + s1=2, + idx_seq=idx_seq, + valid_seq=valid_seq, + softmax_scale=1.0, + head_chunk_size=2, + topk_chunk_size=2, + sk=3, + hn=2, + ) + + assert per_head.shape == (2, 2) + assert not per_head.requires_grad + assert torch.allclose(per_head.sum(dim=-1), torch.tensor([2.0, 2.0]), atol=1e-6) + assert per_head[1, 1] == 0 + + +def test_tilelang_dsa_fused_hook_guard_paths(monkeypatch): + q = torch.empty(2, 1, 2, 4) + k = torch.empty(3, 1, 4) + weights = torch.empty(2, 1, 2) + starts = torch.zeros(2, dtype=torch.int32) + ends = torch.ones(2, dtype=torch.int32) + + monkeypatch.setattr(tilelang_dsa, "lighting_indexer_indices", None) + assert tilelang_dsa.fused_qk_topk_lighting(q, k, weights, 2, starts, ends, 1) is None + + monkeypatch.setattr(tilelang_dsa, "lighting_indexer_indices", lambda *args, **kwargs: None) + assert tilelang_dsa.fused_qk_topk_lighting(q.squeeze(1), k, weights, 2, starts, ends, 1) is None + assert tilelang_dsa.fused_qk_topk_lighting(q, k[:, :0], weights, 2, starts, ends, 1) is None + + monkeypatch.setattr(tilelang_dsa, "lighting_indexer", None) + query = torch.empty(2, 1, 2, 4) + key = torch.empty(3, 1, 1, 4) + assert ( + tilelang_dsa.fused_qk_topk_lighting_with_streaming_sparse_kl( + q, + k, + weights, + 2, + starts, + ends, + 1, + query, + key, + 1.0, + 0.1, + SimpleNamespace(tp=SimpleNamespace(size=lambda: 1)), + ) + is None + ) + + def fail_lighting_indexer(*_args, **_kwargs): + raise AssertionError("unsupported indexer head count should fall back before TileLang") + + monkeypatch.setattr(tilelang_dsa, "lighting_indexer", fail_lighting_indexer) + assert ( + tilelang_dsa.fused_qk_topk_lighting_with_streaming_sparse_kl( + torch.empty(2, 1, 7, 4), + k, + torch.empty(2, 1, 7), + 2, + starts, + ends, + 1, + torch.empty(2, 1, 2, 4), + key, + 1.0, + 0.1, + SimpleNamespace(tp=SimpleNamespace(size=lambda: 1)), + ) + is None + ) + + monkeypatch.setattr(tilelang_dsa, "SparseMLA", None) + topk_indices = torch.zeros(1, 2, 64, dtype=torch.int32) + assert tilelang_dsa.fused_sparse_mla_absorbed(query, key, topk_indices, 1.0, 512) is None + + class FakeSparseMLA: + @staticmethod + def apply(q_t, kv_t, idx_t, softmax_scale): + del q_t, kv_t, idx_t, softmax_scale + return torch.empty(2, 2, 128), torch.empty(2, 2) + + monkeypatch.setattr(tilelang_dsa, "SparseMLA", FakeSparseMLA) + assert ( + tilelang_dsa.fused_sparse_mla_absorbed(query.squeeze(1), key, topk_indices, 1.0, 512) + is None + ) + assert ( + tilelang_dsa.fused_sparse_mla_absorbed(query, key.squeeze(2), topk_indices, 1.0, 512) + is None + ) + assert tilelang_dsa.fused_sparse_mla_absorbed(query, key, topk_indices[:0], 1.0, 512) is None + assert tilelang_dsa.fused_sparse_mla_absorbed(query, key, topk_indices[:, :1], 1.0, 512) is None + assert ( + tilelang_dsa.fused_sparse_mla_absorbed(query, key[..., :3], topk_indices, 1.0, 512) is None + ) + assert ( + tilelang_dsa.fused_sparse_mla_absorbed(query, key, topk_indices[..., :63], 1.0, 512) is None + ) + + query_supported = torch.empty(2, 1, 3, 576) + key_supported = torch.empty(2, 1, 1, 576) + topk_supported = torch.zeros(1, 2, 64, dtype=torch.int32) + assert ( + tilelang_dsa.fused_sparse_mla_absorbed( + query_supported, key_supported, topk_supported, 1.0, 256 + ) + is None + ) + assert ( + tilelang_dsa.fused_sparse_mla_absorbed( + query_supported, key_supported, topk_supported[..., :63], 1.0, 512 + ) + is None + ) + + class FailSparseMLA: + @staticmethod + def apply(*_args, **_kwargs): + raise AssertionError( + "unsupported SparseMLA head count should fall back before TileLang" + ) + + monkeypatch.setattr(tilelang_dsa, "SparseMLA", FailSparseMLA) + assert ( + tilelang_dsa.fused_sparse_mla_absorbed( + torch.empty(2, 1, 96, 576), key_supported, topk_supported, 1.0, 512 + ) + is None + ) + + monkeypatch.setattr(tilelang_dsa, "SparseMLA", FakeSparseMLA) + assert ( + tilelang_dsa.fused_sparse_mla_absorbed( + query_supported, key_supported, topk_supported, 1.0, 512 + ) + is None + ) + + +def test_fused_qk_topk_lighting_sanitizes_mocked_tilelang_indices(monkeypatch): + q = torch.empty(3, 1, 2, 4, dtype=torch.bfloat16) + k = torch.empty(5, 1, 4, dtype=torch.bfloat16) + weights = torch.empty(3, 1, 2) + starts = torch.tensor([0, 2, 4], dtype=torch.int32) + ends = torch.tensor([2, 4, 5], dtype=torch.int32) + calls = [] + + def fake_lighting_indexer_indices( + index_q, index_k, index_w, starts_arg, ends_arg, index_topk, use_relu=True + ): + del index_k, index_w, index_topk + calls.append((tuple(index_q.shape), starts_arg.clone(), ends_arg.clone(), use_relu)) + return torch.stack((starts_arg, ends_arg), dim=-1).to(torch.int32) + + monkeypatch.setattr(tilelang_dsa, "lighting_indexer_indices", fake_lighting_indexer_indices) + + topk = tilelang_dsa.fused_qk_topk_lighting( + q, k, weights, index_topk=2, starts=starts, ends=ends, block_size=2, use_relu=False + ) + + assert torch.equal(topk, torch.tensor([[[0, -1], [2, -1], [4, -1]]], dtype=torch.int32)) + assert [call[0] for call in calls] == [(2, 2, 4), (1, 2, 4)] + assert all(call[3] is False for call in calls) + + +def test_fused_sparse_mla_absorbed_batches_mocked_tilelang_outputs(monkeypatch): + class FakeSparseMLA: + @staticmethod + def apply(q_t, kv_t, idx_t, softmax_scale): + assert q_t.shape == (2, 2, 16, 576) + assert kv_t.shape == (2, 2, 1, 576) + assert idx_t.shape == (2, 2, 1, 64) + assert softmax_scale == 0.25 + batch_sums = q_t.float().sum(dim=(1, 2, 3)).to(dtype=q_t.dtype) + out = batch_sums.view(q_t.size(0), 1, 1, 1).expand( + q_t.size(0), q_t.size(1), q_t.size(2), 512 + ) + lse = torch.zeros(q_t.size(0), q_t.size(1), q_t.size(2)) + return out, lse + + monkeypatch.setattr(tilelang_dsa, "SparseMLA", FakeSparseMLA) + query = torch.zeros(2, 2, 16, 576, dtype=torch.bfloat16) + query[:, 1].fill_(1.0) + key = torch.zeros(2, 2, 1, 576, dtype=torch.bfloat16) + topk_indices = torch.zeros(2, 2, 64, dtype=torch.int32) + + output = tilelang_dsa.fused_sparse_mla_absorbed( + query, key, topk_indices, softmax_scale=0.25, v_channels=512 + ) + + assert output.shape == (2, 2, 16, 512) + assert torch.equal(output[:, 0], torch.zeros_like(output[:, 0])) + assert torch.equal(output[:, 1], torch.full_like(output[:, 1], 18432.0)) + + +def test_fused_sparse_mla_absorbed_pads_small_head_count_without_gradient_leak(monkeypatch): + class FakeSparseMLA: + @staticmethod + def apply(q_t, kv_t, idx_t, softmax_scale): + assert q_t.shape == (1, 2, 16, 576) + assert kv_t.shape == (1, 2, 1, 576) + assert idx_t.shape == (1, 2, 1, 64) + assert softmax_scale == 0.25 + assert torch.count_nonzero(q_t[:, :, 8:]) == 0 + out = q_t[..., :512] + kv_t[..., :512] + return out, torch.zeros(q_t.shape[:-1], dtype=torch.float32) + + monkeypatch.setattr(tilelang_dsa, "SparseMLA", FakeSparseMLA) + query = torch.randn(2, 1, 8, 576, dtype=torch.bfloat16, requires_grad=True) + key = torch.randn(2, 1, 1, 576, dtype=torch.bfloat16, requires_grad=True) + topk_indices = torch.zeros(1, 2, 64, dtype=torch.int32) + + output = tilelang_dsa.fused_sparse_mla_absorbed( + query, key, topk_indices, softmax_scale=0.25, v_channels=512 + ) + + assert output is not None + assert output.shape == (2, 1, 8, 512) + output.float().sum().backward() + assert torch.equal(query.grad[..., :512], torch.ones_like(query.grad[..., :512])) + assert torch.count_nonzero(query.grad[..., 512:]) == 0 + assert torch.equal(key.grad[..., :512], torch.full_like(key.grad[..., :512], 8.0)) + assert torch.count_nonzero(key.grad[..., 512:]) == 0 + + +def test_streaming_sparse_kl_path_with_mocked_tilelang_indexer(monkeypatch): + q = torch.empty(2, 1, 2, 4, dtype=torch.bfloat16) + k = torch.empty(4, 1, 4, dtype=torch.bfloat16) + weights = torch.empty(2, 1, 2) + starts = torch.tensor([0, 0], dtype=torch.int32) + ends = torch.tensor([4, 4], dtype=torch.int32) + query = torch.empty(2, 1, 2, 4, dtype=torch.bfloat16) + key = torch.empty(4, 1, 1, 4, dtype=torch.bfloat16) + query_valid_rows = torch.tensor([[True, False]]) + + def fake_lighting_indexer( + index_q, + index_k, + index_w, + starts_arg, + ends_arg, + index_topk, + topk_indices=None, + use_relu=True, + ): + del index_k, index_w, starts_arg, ends_arg, topk_indices, use_relu + topk_scores = torch.zeros(index_q.size(0), index_topk) + topk = torch.tensor([[0, 1], [2, 3]], dtype=torch.int32)[: index_q.size(0)] + return topk_scores, topk + + def fake_compute_topk_target_chunk_sum(**kwargs): + idx_seq = kwargs["idx_seq"] + return torch.ones(idx_seq.shape, dtype=torch.float32, device=idx_seq.device) + + monkeypatch.setattr(tilelang_dsa, "lighting_indexer", fake_lighting_indexer) + monkeypatch.setattr(tilelang_dsa, "is_supported_indexer_bwd_head_count", lambda *_args: True) + monkeypatch.setattr( + tilelang_dsa, "_compute_topk_target_chunk_sum", fake_compute_topk_target_chunk_sum + ) + + topk, loss = tilelang_dsa.fused_qk_topk_lighting_with_streaming_sparse_kl( + q=q, + k=k, + weights=weights, + index_topk=2, + starts=starts, + ends=ends, + block_size=2, + query=query, + key=key, + softmax_scale=0.5, + loss_coeff=2.0, + pg_collection=SimpleNamespace(tp=SimpleNamespace(size=lambda: 1)), + query_valid_rows=query_valid_rows, + calculate_per_token_loss=False, + seq_chunk_size=1, + head_chunk_size=1, + topk_chunk_size=1, + use_relu=False, + ) + + assert torch.equal(topk, torch.tensor([[[0, 1], [2, 3]]], dtype=torch.int32)) + assert loss.item() == 0.0 + + +def test_streaming_sparse_kl_uses_fused_target_when_supported(monkeypatch): + q = torch.empty(2, 1, 2, 4, dtype=torch.bfloat16) + k = torch.empty(4, 1, 4, dtype=torch.bfloat16) + weights = torch.empty(2, 1, 2) + starts = torch.zeros(2, dtype=torch.int32) + ends = torch.full((2,), 4, dtype=torch.int32) + query = torch.empty(2, 1, 2, 4, dtype=torch.bfloat16) + key = torch.empty(4, 1, 1, 4, dtype=torch.bfloat16) + calls = [] + + def fake_lighting_indexer( + index_q, + index_k, + index_w, + starts_arg, + ends_arg, + index_topk, + topk_indices=None, + use_relu=True, + ): + del index_k, index_w, starts_arg, ends_arg, topk_indices, use_relu + topk_scores = torch.zeros(index_q.size(0), index_topk, requires_grad=True) + topk = torch.tensor([[0, 1], [2, 3]], dtype=torch.int32)[: index_q.size(0)] + return topk_scores, topk + + def fake_target(query_arg, key_arg, indices_arg, softmax_scale): + calls.append((query_arg, key_arg, indices_arg.clone(), softmax_scale)) + return torch.ones(indices_arg.shape, dtype=torch.float32) + + def fail_python_target(**_kwargs): + raise AssertionError("the PyTorch target path should not run") + + monkeypatch.setattr(tilelang_dsa, "lighting_indexer", fake_lighting_indexer) + monkeypatch.setattr(tilelang_dsa, "is_supported_indexer_bwd_head_count", lambda *_args: True) + monkeypatch.setattr(tilelang_dsa, "_can_use_fused_sparse_indexer_target", lambda *_args: True) + monkeypatch.setattr(tilelang_dsa, "sparse_indexer_target_interface", fake_target) + monkeypatch.setattr(tilelang_dsa, "_can_use_fused_sparse_indexer_kl", lambda *_args: False) + monkeypatch.setattr(tilelang_dsa, "_compute_topk_target_chunk_sum", fail_python_target) + + topk, loss = tilelang_dsa.fused_qk_topk_lighting_with_streaming_sparse_kl( + q=q, + k=k, + weights=weights, + index_topk=2, + starts=starts, + ends=ends, + block_size=2, + query=query, + key=key, + softmax_scale=0.5, + loss_coeff=2.0, + pg_collection=SimpleNamespace(tp=SimpleNamespace(size=lambda: 1)), + seq_chunk_size=2, + ) + + assert torch.equal(topk, torch.tensor([[[0, 1], [2, 3]]], dtype=torch.int32)) + assert loss.item() == 0.0 + assert len(calls) == 1 + assert calls[0][0].shape == (2, 2, 4) + assert calls[0][1].shape == (4, 4) + assert calls[0][3] == 0.5 + + +@pytest.mark.parametrize("heads", [48, 96]) +def test_fused_sparse_indexer_target_and_kl_match_reference(heads): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for TileLang indexer-loss tests") + if not tilelang_indexer_loss.HAVE_TILELANG: + pytest.skip("TileLang indexer-loss kernels are unavailable") + + torch.manual_seed(1234 + heads) + seq_len = 2 + key_len = 256 + topk = 256 + dim = 576 + softmax_scale = dim**-0.5 + query = torch.randn(seq_len, heads, dim, device="cuda", dtype=torch.bfloat16) + key = torch.randn(key_len, dim, device="cuda", dtype=torch.bfloat16) + topk_indices = torch.arange(topk, device="cuda", dtype=torch.int32).repeat(seq_len, 1) + topk_indices[1, -16:] = -1 + valid = topk_indices >= 0 + + target = tilelang_indexer_loss.sparse_indexer_target_interface( + query, key, topk_indices, softmax_scale + ) + safe_indices = topk_indices.clamp(min=0).to(torch.int64) + selected_key = key.index_select(0, safe_indices.reshape(-1)).view(seq_len, topk, dim) + reference_scores = ( + torch.einsum("shd,skd->shk", query.float(), selected_key.float()) * softmax_scale + ) + reference_scores = reference_scores.masked_fill(~valid.unsqueeze(1), float("-inf")) + reference_target = torch.softmax(reference_scores, dim=-1).masked_fill(~valid.unsqueeze(1), 0.0) + reference_target = reference_target.sum(dim=1) + torch.testing.assert_close(target, reference_target, rtol=2e-2, atol=2e-2) + + logits = torch.randn(seq_len, topk, device="cuda", dtype=torch.float32, requires_grad=True) + loss = tilelang_indexer_loss.SparseIndexerKLLoss.apply(target, logits, valid) + loss.backward() + + normalized_target = dsa_indexer_loss.normalize_indexer_target(reference_target) + reference_log_probs = dsa_masking.masked_log_softmax(logits.detach(), valid, dim=-1) + reference_loss = dsa_indexer_loss.indexer_kl_sum(normalized_target, reference_log_probs, valid) + reference_grad = ( + reference_log_probs.exp().masked_fill(~valid, 0.0) - normalized_target + ).masked_fill(~valid, 0.0) + torch.testing.assert_close(loss, reference_loss, rtol=2e-3, atol=2e-3) + torch.testing.assert_close(logits.grad, reference_grad, rtol=2e-3, atol=2e-3) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.float32]) +def test_tilelang_ops_decline_non_bfloat16_inputs(monkeypatch, dtype): + def fail_if_called(*_args, **_kwargs): + raise AssertionError("TileLang kernel should not run for non-BF16 inputs") + + class FailSparseMLA: + @staticmethod + def apply(*args): + return fail_if_called(*args) + + monkeypatch.setattr(tilelang_dsa, "lighting_indexer_indices", fail_if_called) + monkeypatch.setattr(tilelang_dsa, "lighting_indexer", fail_if_called) + monkeypatch.setattr(tilelang_dsa, "SparseMLA", FailSparseMLA) + + q_indexer = torch.zeros((1, 1, 1, 1), dtype=dtype) + k_indexer = torch.zeros((1, 1, 1), dtype=dtype) + weights = torch.zeros((1, 1, 1), dtype=dtype) + starts = torch.tensor([0], dtype=torch.int32) + ends = torch.tensor([1], dtype=torch.int32) + query = torch.zeros((1, 1, 1, 1), dtype=dtype) + key = torch.zeros((1, 1, 1, 1), dtype=dtype) + topk_indices = torch.zeros((1, 1, 1), dtype=torch.int32) + + assert ( + tilelang_dsa.fused_qk_topk_lighting(q_indexer, k_indexer, weights, 1, starts, ends, 128) + is None + ) + assert ( + tilelang_dsa.fused_qk_topk_lighting_with_streaming_sparse_kl( + q=q_indexer, + k=k_indexer, + weights=weights, + index_topk=1, + starts=starts, + ends=ends, + block_size=128, + query=query, + key=key, + softmax_scale=1.0, + loss_coeff=0.01, + pg_collection=object(), + ) + is None + ) + assert tilelang_dsa.fused_sparse_mla_absorbed(query, key, topk_indices, 1.0, 1) is None + + +@pytest.mark.parametrize("num_heads", [8, 64]) +def test_fused_sparse_mla_absorbed_accepts_thd_sentinels(num_heads): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for TileLang SparseMLA tests") + if tilelang_dsa.SparseMLA is None: + pytest.skip("TileLang SparseMLA kernel is unavailable") + + torch.manual_seed(1234) + torch.cuda.manual_seed(1234) + + # Match the sequence bucket so the kernel receives the original B=1 tensor views. + # This exercises the canonical batch-stride handling rather than hiding it with padding. + seqlen = 256 + dim = 576 + v_channels = 512 + topk = 64 + + query = torch.randn( + (seqlen, 1, num_heads, dim), dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + key = torch.randn((seqlen, 1, 1, dim), dtype=torch.bfloat16, device="cuda", requires_grad=True) + topk_indices = torch.full((1, seqlen, topk), -1, dtype=torch.int32, device="cuda") + for row in range(1, seqlen): + valid = min(row, topk) + topk_indices[0, row, :valid] = torch.arange(valid, dtype=torch.int32, device="cuda") + + output = tilelang_dsa.fused_sparse_mla_absorbed( + query, key, topk_indices, softmax_scale=1.0 / math.sqrt(dim), v_channels=v_channels + ) + + assert output is not None + assert output.shape == (seqlen, 1, num_heads, v_channels) + assert torch.isfinite(output).all() + assert output[0].abs().max() == 0 + + output.float().square().mean().backward() + assert query.grad is not None + assert key.grad is not None + assert torch.isfinite(query.grad).all() + assert torch.isfinite(key.grad).all() + assert query.grad[0].abs().max() == 0