diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index be7a2fa781..79a9972901 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/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) diff --git a/unit_tests/models/glm5_next/test_token_major_mla_bmm.py b/unit_tests/models/glm5_next/test_token_major_mla_bmm.py new file mode 100644 index 0000000000..5de243d770 --- /dev/null +++ b/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()