perf(glm): write MLA projections directly in token-major layout

lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py

@@ -152,7 +152,10 @@ def _get_qkv(self, input, infer_state, layer_weight):
         return q, cache_kv
 
     def _context_attention_kernel(self, q, kv, infer_state, layer_weight, out=None):
-        q = layer_weight.k_b_proj_.bmm(q.transpose(0, 1)).transpose(0, 1).contiguous()
+        q_nope = self.alloc_tensor((q.shape[0], self.tp_q_head_num_, self.kv_lora_rank), dtype=q.dtype, device=q.device)
+        # Write [tokens, heads, latent] directly for sparse prefill attention.
+        layer_weight.k_b_proj_.bmm(q.transpose(0, 1), out=q_nope.transpose(0, 1))
+        q = q_nope
         topk_mem_indices, topk_indices = self.indexer._get_indices(
             hidden_states=infer_state.get_topk_indices_params["hidden_states"],
             q_lora=infer_state.get_topk_indices_params["q_lora"],
@@ -273,9 +276,12 @@ def _token_attention_kernel(self, q, infer_state, layer_weight, out=None):
     def _get_o(self, input, infer_state, layer_weight):
         if infer_state.need_dp_prefill_balance:
             input = infer_state._all_to_all_balance_get(data=input)
-        # Both sparse prefill and decode return attention in MLA latent space.
-        input = layer_weight.v_b_proj_.bmm(input.transpose(0, 1)).transpose(0, 1)
-        output = layer_weight.o_weight_.mm(input.reshape(-1, self.tp_q_head_num_ * self.v_head_dim))
+        values = self.alloc_tensor(
+            (input.shape[0], self.tp_q_head_num_, self.v_head_dim), dtype=input.dtype, device=input.device
+        )
+        # Token-major output lets prefill and decode flatten the heads without copying.
+        layer_weight.v_b_proj_.bmm(input.transpose(0, 1), out=values.transpose(0, 1))
+        output = layer_weight.o_weight_.mm(values.view(-1, self.tp_q_head_num_ * self.v_head_dim))
         return self._tpsp_reduce(input=output, infer_state=infer_state)
 
     def _kda_projections(self, hidden_states, infer_state, layer_weight):

unit_tests/models/glm5_next/test_layers.py

@@ -113,7 +113,8 @@ def test_kda_projections_need_no_sparse_attention_config(config):
 
 
 @pytest.mark.parametrize("rank", [0, 1])
-def test_mla_loading_preserves_replicated_qkv_and_sharded_bmm(config, monkeypatch, rank):
+@pytest.mark.parametrize("is_prefill", [False, True])
+def test_mla_loading_preserves_replicated_qkv_and_sharded_bmm(config, monkeypatch, rank, is_prefill):
     monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "2")
     monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", str(rank))
     layer = Glm5NextTransformerLayerInfer(3, config)
@@ -133,7 +134,7 @@ def test_mla_loading_preserves_replicated_qkv_and_sharded_bmm(config, monkeypatc
         {f"model.language_model.layers.3.self_attn.{name}": value for name, value in tensors.items()}
     )
     hidden = random_weight(3, 128)
-    state = SimpleNamespace(need_dp_prefill_balance=False)
+    state = SimpleNamespace(is_prefill=is_prefill, need_dp_prefill_balance=False)
     q, kv = layer._get_qkv(hidden, state, weight)
     q_lora = rmsnorm(F.linear(hidden, q_a), config["rms_norm_eps"])
     expected_q = F.linear(q_lora, q_b).view(3, 4, 128)[:, rank * 2 : (rank + 1) * 2]

unit_tests/models/glm5_next/test_token_major_mla_bmm.py

@@ -0,0 +1,92 @@
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import BMMWeightTpl
+from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer
+
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+def make_projection(layer, heads: int, inner: int, cols: int, dtype: torch.dtype):
+    projection = BMMWeightTpl.__new__(BMMWeightTpl)
+    projection.weight = torch.randn(heads, inner, cols, device="cuda", dtype=dtype) * 0.1
+    layer.tp_q_head_num_ = heads
+    if inner == 256:
+        layer.kv_lora_rank = cols
+        layer.layer_num_ = 0
+        layer.softmax_scale = inner ** -0.5
+        layer.indexer = SimpleNamespace(_get_indices=lambda **kwargs: (None, None))
+        state = SimpleNamespace(
+            att_state_role="main",
+            att_state_manager=Mock(),
+            prefill_att_state=SimpleNamespace(prefill_att=lambda **kwargs: kwargs["q"]),
+        )
+        weight = SimpleNamespace(k_b_proj_=projection)
+
+        def run(x):
+            state.get_topk_indices_params = {"hidden_states": None, "q_lora": None}
+            return layer._context_attention_kernel(x, None, state, weight)
+
+    else:
+        layer.v_head_dim = cols
+        layer._tpsp_reduce = lambda input, infer_state: input
+        state = SimpleNamespace(is_prefill=False, need_dp_prefill_balance=False)
+        weight = SimpleNamespace(v_b_proj_=projection, o_weight_=SimpleNamespace(mm=lambda x: x))
+
+        def run(x):
+            return layer._get_o(x, state, weight).view(x.shape[0], heads, cols)
+
+    return projection, run
+
+
+@pytest.mark.parametrize("rows", [1, 128, 1024, 1659, 4096, 8192])
+@pytest.mark.parametrize("heads", [16, 64])
+@pytest.mark.parametrize("inner,cols", [(256, 512), (512, 256)])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_token_major_projection_matches_reference_and_keeps_guard_rows(rows, heads, inner, cols, dtype):
+    torch.manual_seed(272)
+    layer = Glm5NextTransformerLayerInfer.__new__(Glm5NextTransformerLayerInfer)
+    projection, run = make_projection(layer, heads, inner, cols, dtype)
+    x = torch.randn(rows, heads, inner, device="cuda", dtype=dtype)
+    expected = projection.bmm(x.transpose(0, 1)).transpose(0, 1).contiguous()
+    slab = torch.full((rows + 2, heads, cols), 113.0, device="cuda", dtype=dtype)
+    output = slab[1:-1]
+    layer.alloc_tensor = Mock(return_value=output)
+    actual = run(x)
+    assert actual.data_ptr() == output.data_ptr() and actual.is_contiguous()
+    assert torch.equal(actual, expected)
+    assert (slab[0] == 113).all() and (slab[-1] == 113).all()
+
+
+@pytest.mark.parametrize("inner,cols", [(256, 512), (512, 256)])
+def test_token_major_projection_accepts_head_padding(inner, cols):
+    # TP4 FlashMLA crops a 64-head output to 16 heads without copying it.
+    layer = Glm5NextTransformerLayerInfer.__new__(Glm5NextTransformerLayerInfer)
+    projection, run = make_projection(layer, 16, inner, cols, torch.bfloat16)
+    x = torch.randn(128, 64, inner, device="cuda", dtype=torch.bfloat16)[:, :16]
+    assert not x.is_contiguous()
+    expected = projection.bmm(x.transpose(0, 1)).transpose(0, 1)
+    actual = run(x)
+    assert actual.is_contiguous()
+    assert torch.equal(actual, expected)
+
+
+@pytest.mark.parametrize("inner,cols", [(256, 512), (512, 256)])
+def test_token_major_projection_replays_new_inputs_in_cuda_graph(inner, cols):
+    layer = Glm5NextTransformerLayerInfer.__new__(Glm5NextTransformerLayerInfer)
+    projection, run = make_projection(layer, 16, inner, cols, torch.bfloat16)
+    x = torch.randn(128, 16, inner, device="cuda", dtype=torch.bfloat16)
+    for _ in range(3):
+        run(x)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        actual = run(x)
+    x.normal_()
+    graph.replay()
+    expected = projection.bmm(x.transpose(0, 1)).transpose(0, 1)
+    assert torch.equal(actual, expected)