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

lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py

@@ -6,8 +6,9 @@
 
 import torch
 
-from lightllm.common.basemodel import TransformerLayerInferTpl
+from lightllm.common.basemodel import InferStateInfo, TransformerLayerInferTpl
 from lightllm.common.basemodel.attention.base_att import AttControl
+from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import BMMWeightTpl
 from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward
 from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import (
     silu_and_mul_fwd,
@@ -20,6 +21,7 @@
 from lightllm.common.triton_utils.autotuner import Autotuner
 from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer
 from lightllm.utils.envs_utils import get_env_start_args
+from lightllm.utils.device_utils import is_hopper
 from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor
 
 
@@ -61,6 +63,19 @@ def __init__(self, layer_num, network_config):
             self.kv_lora_rank = network_config["kv_lora_rank"]
             self.v_head_dim = network_config["v_head_dim"]
             self.softmax_scale = self.qk_nope_head_dim ** -0.5
+            args = get_env_start_args()
+            self.prefill_token_major_mla_bmm = (
+                args.run_mode == "prefill"
+                and args.ssm_state_mode == "native"
+                and self.tp_world_size_ == 1
+                and args.enable_ep_moe
+                and not args.enable_prefill_cudagraph
+                and not args.enable_prefill_microbatch_overlap
+                and not args.enable_tpsp_mix_mode
+                and (self.tp_q_head_num_, self.qk_nope_head_dim, self.kv_lora_rank, self.v_head_dim)
+                == (64, 256, 512, 256)
+                and is_hopper()
+            )
             self.indexer = Glm5NextNsaInfer(
                 layer_idx=self.layer_num_,
                 network_config=self.network_config_,
@@ -151,8 +166,33 @@ def _get_qkv(self, input, infer_state, layer_weight):
         )
         return q, cache_kv
 
+    def _sparse_mla_bmm(
+        self, x: torch.Tensor, projection: BMMWeightTpl, infer_state: InferStateInfo, out_features: int
+    ) -> torch.Tensor:
+        # Preserve the reference arithmetic for short, entirely cold waves.
+        if (
+            self.prefill_token_major_mla_bmm
+            and infer_state.is_prefill
+            and not infer_state.is_cuda_graph_capture
+            and not infer_state.need_dp_prefill_balance
+            and not (infer_state.max_cache_len == 0 and 0 < infer_state.max_kv_seq_len <= 2048)
+            and x.is_cuda
+            and x.dtype == projection.weight.dtype == torch.bfloat16
+            and 1024 <= x.shape[0] <= 8192
+            and x.is_contiguous()
+            and projection.weight.is_contiguous()
+            and projection.weight.shape == (x.shape[1], x.shape[2], out_features)
+        ):
+            # Write directly in the layout consumed by attention/the output projection.
+            output = self.alloc_tensor(shape=(x.shape[0], x.shape[1], out_features), dtype=x.dtype, device=x.device)
+            projection.bmm(input_tensor=x.transpose(0, 1), out=output.transpose(0, 1))
+            return output
+        return projection.bmm(input_tensor=x.transpose(0, 1)).transpose(0, 1)
+
     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 = self._sparse_mla_bmm(
+            x=q, projection=layer_weight.k_b_proj_, infer_state=infer_state, out_features=self.kv_lora_rank
+        ).contiguous()
         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"],
@@ -274,7 +314,9 @@ 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)
+        input = self._sparse_mla_bmm(
+            x=input, projection=layer_weight.v_b_proj_, infer_state=infer_state, out_features=self.v_head_dim
+        )
         output = layer_weight.o_weight_.mm(input.reshape(-1, self.tp_q_head_num_ * self.v_head_dim))
         return self._tpsp_reduce(input=output, infer_state=infer_state)
 

unit_tests/models/glm5_next/test_token_major_mla_bmm.py

@@ -0,0 +1,69 @@
+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_layer_and_state() -> tuple[Glm5NextTransformerLayerInfer, SimpleNamespace]:
+    layer = Glm5NextTransformerLayerInfer.__new__(Glm5NextTransformerLayerInfer)
+    layer.prefill_token_major_mla_bmm = True
+    infer = SimpleNamespace(
+        is_prefill=True,
+        is_cuda_graph_capture=False,
+        need_dp_prefill_balance=False,
+        max_cache_len=131072,
+        max_kv_seq_len=139264,
+    )
+    return layer, infer
+
+
+@pytest.mark.parametrize("rows", [1024, 1659, 4096, 8192])
+@pytest.mark.parametrize("inner,cols", [(256, 512), (512, 256)])
+def test_token_major_projection_matches_reference_and_keeps_guard_rows(rows: int, inner: int, cols: int) -> None:
+    torch.manual_seed(272)
+    layer, infer = make_layer_and_state()
+    projection = BMMWeightTpl.__new__(BMMWeightTpl)
+    projection.weight = torch.randn(64, inner, cols, device="cuda", dtype=torch.bfloat16) * 0.1
+    x = torch.randn(rows, 64, inner, device="cuda", dtype=torch.bfloat16)
+    expected = projection.bmm(x.transpose(0, 1)).transpose(0, 1).contiguous()
+    slab = torch.full((rows + 2, 64, cols), 113.0, device="cuda", dtype=torch.bfloat16)
+    output = slab[1:-1]
+    layer.alloc_tensor = Mock(return_value=output)
+    actual = layer._sparse_mla_bmm(x=x, projection=projection, infer_state=infer, out_features=cols)
+    assert actual is output and actual.is_contiguous()
+    assert torch.equal(actual, expected)
+    assert (slab[0] == 113).all() and (slab[-1] == 113).all()
+    layer.alloc_tensor.assert_called_once_with(shape=(rows, 64, cols), dtype=x.dtype, device=x.device)
+
+
+@pytest.mark.parametrize("case", ["decode", "capture", "balance", "cold", "small", "disabled", "strided"])
+def test_unsupported_execution_keeps_reference_bmm(case: str) -> None:
+    layer, infer = make_layer_and_state()
+    if case == "decode":
+        infer.is_prefill = False
+    elif case == "capture":
+        infer.is_cuda_graph_capture = True
+    elif case == "balance":
+        infer.need_dp_prefill_balance = True
+    elif case == "cold":
+        infer.max_cache_len, infer.max_kv_seq_len = 0, 1024
+    elif case == "disabled":
+        layer.prefill_token_major_mla_bmm = False
+    rows = 128 if case == "small" else 1024
+    x = torch.randn(rows, 64, 257 if case == "strided" else 256, device="cuda", dtype=torch.bfloat16)[..., :256]
+    projection = BMMWeightTpl.__new__(BMMWeightTpl)
+    projection.weight = torch.randn(64, 256, 512, device="cuda", dtype=torch.bfloat16)
+    reference = projection.bmm(x.transpose(0, 1))
+    projection.bmm = Mock(return_value=reference)
+    layer.alloc_tensor = Mock(side_effect=AssertionError("unsupported path must not use direct output"))
+    actual = layer._sparse_mla_bmm(x=x, projection=projection, infer_state=infer, out_features=512)
+    assert torch.equal(actual, reference.transpose(0, 1))
+    assert "out" not in projection.bmm.call_args.kwargs
+    layer.alloc_tensor.assert_not_called()