[Perf] Add CuTe DSL DSA indexer

pyproject.toml

@@ -67,6 +67,9 @@ path = "xtuner/version.py"
 allow-direct-references = true
 
 [project.optional-dependencies]
+cute = [
+  "nvidia-cutlass-dsl==4.5.2",
+]
 rl = [
   "ray[default]",
   "httpx",

tests/engine/test_glm52_moe_train_engine.py

@@ -87,6 +87,7 @@ def _tiny_checkpoint_config(dispatcher: str | None, ep_size: int) -> Glm52MoECon
             index_head_dim=4,
             index_n_heads=2,
             indexer_types=["full", "shared"],
+            indexer_backend="torch",
             sparse_mla_backend="torch",
         ),
         hf_head_dim=4,

tests/model/test_glm52_moe.py

@@ -60,6 +60,7 @@ def _tiny_glm52_config() -> Glm52MoEConfig:
             index_head_dim=4,
             index_n_heads=2,
             indexer_types=["full", "shared", "shared", "full"],
+            indexer_backend="torch",
             sparse_mla_backend="torch",
         ),
         hf_head_dim=4,

tests/model/test_glm52_mtp_checkpoint_repro.py

@@ -49,6 +49,7 @@ def _tiny_mtp_config(ep_size: int, mtp_num_layers: int, compile_model: bool) ->
             index_head_dim=4,
             index_n_heads=2,
             indexer_types=["full", "shared", "full"],
+            indexer_backend="torch",
         ),
         hf_head_dim=4,
         qk_head_dim=8,

tests/module/attention/test_dsa_mla.py

@@ -3,6 +3,7 @@
 TestTorchSparseMLA
     test_padded_indices_support_int32_and_backward: PyTorch 后端处理 padding、int32 和反向传播。
 TestDSAAttention
+    test_indexer_backend_is_independent_from_sparse_mla_backend: indexer backend 独立配置且默认 TileLang。
     test_packed_inputs_respect_causal_boundaries_and_backward: packed attention 遵守分段因果边界并可反传。
     test_shared_layers_reuse_topk_without_cross_context_leak: shared layer 复用当前样本 top-k 且不跨样本泄漏。
     test_reentrant_checkpoint_reuses_and_releases_topk: checkpoint 重算复用并最终释放 top-k。
@@ -115,6 +116,7 @@ def _tiny_dsa_attention(
         index_head_dim=4,
         index_n_heads=2,
         indexer_types=indexer_types,
+        indexer_backend="torch",
         sparse_mla_backend="torch",
     ).build(hidden_size=4, layer_idx=layer_idx)
 
@@ -167,6 +169,29 @@ def test_padded_indices_support_int32_and_backward(self):
 
 
 class TestDSAAttention:
+    def test_indexer_backend_is_independent_from_sparse_mla_backend(self):
+        config = DSAMLAConfig(
+            num_attention_heads=2,
+            head_dim=2,
+            kv_lora_rank=3,
+            q_lora_rank=4,
+            qk_nope_head_dim=2,
+            qk_rope_head_dim=2,
+            v_head_dim=3,
+            index_topk=4,
+            index_head_dim=4,
+            index_n_heads=2,
+            sparse_mla_backend="tilelang",
+        )
+        assert config.indexer_backend == "tilelang"
+
+        config.indexer_backend = "torch"
+        attention = config.build(hidden_size=4)
+
+        assert attention.indexer_backend == "torch"
+        assert attention.sparse_mla_backend == "tilelang"
+        assert attention.indexer.indexer_backend == "torch"
+
     def test_packed_inputs_respect_causal_boundaries_and_backward(self):
         # 验证 packed attention 不跨子序列取 key，并能对真实输入完成有限反向传播。
         torch.manual_seed(0)

tests/module/test_dense_decoder_layer.py

@@ -30,6 +30,7 @@ def _build_dense_dsa_layer() -> DenseDecoderLayer:
             index_head_dim=4,
             index_n_heads=2,
             indexer_types=["full"],
+            indexer_backend="torch",
             sparse_mla_backend="torch",
         ),
     )