From 4f87e08e5710b37958096188104230b0cd35db98 Mon Sep 17 00:00:00 2001 From: HappyAngel Date: Thu, 13 Aug 2026 21:03:18 +0800 Subject: [PATCH] fix: fix lite csa.py code doesn't match megatron-lm dev rescontruct code --- .../megatron/lite/primitive/modules/attention/csa.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/experimental/lite/megatron/lite/primitive/modules/attention/csa.py b/experimental/lite/megatron/lite/primitive/modules/attention/csa.py index 63b7707e114..7d45faad09d 100644 --- a/experimental/lite/megatron/lite/primitive/modules/attention/csa.py +++ b/experimental/lite/megatron/lite/primitive/modules/attention/csa.py @@ -8,11 +8,11 @@ # Zero-copy imports of the DSv4 THD-CP helpers that live in Megatron Core. The # lite CSA module reuses Core's differentiable kernels, CP row-mapping utilities, # and CuTeDSL layout kernels rather than vendoring them; see the module docstring -# of ``csa_cp_utils`` / ``csa_cp_layout_kernels`` for the exact contracts. +# of ``cp_utils`` / ``cp_layout_kernels`` for the exact contracts. from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region -from megatron.core.transformer.experimental_attention_variant import ( - csa_cp_layout_kernels, - csa_cp_utils as cp_utils, +from megatron.core.transformer.experimental_attention_variant.csa_utils import ( + cp_layout_kernels, + cp_utils, ) from megatron.core.transformer.experimental_attention_variant.csa import ( _unfused_indexer_sparse_attn_from_topk, @@ -969,7 +969,7 @@ def _forward_thd_cp( else (max_seqlen_q // ratio if ratio > 1 else 0) ) topk_idxs, topk_length, indexer_topk_rank_major = ( - csa_cp_layout_kernels.build_attention_indices( + cp_layout_kernels.build_attention_indices( cu_seqlens, global_start, l_local,