From 080a826329a06524b46bfef7b4d6d59757ba8eee Mon Sep 17 00:00:00 2001 From: baishihao Date: Thu, 12 Nov 2026 17:59:11 +0800 Subject: [PATCH 1/2] perf(glm): write prefill MLA projections in token-major layout --- .../layer_infer/transformer_layer_infer.py | 48 ++++++++++++- .../glm5_next/test_token_major_mla_bmm.py | 69 +++++++++++++++++++ 2 files changed, 114 insertions(+), 3 deletions(-) create mode 100644 unit_tests/models/glm5_next/test_token_major_mla_bmm.py 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() From fc845747aa65112fb79560b46b187a9ad060fea4 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Thu, 12 Nov 2026 14:05:37 +0000 Subject: [PATCH 2/2] perf(glm): simplify token-major projections and cover decode --- .../layer_infer/transformer_layer_infer.py | 56 ++------ unit_tests/models/glm5_next/test_layers.py | 5 +- .../glm5_next/test_token_major_mla_bmm.py | 121 +++++++++++------- 3 files changed, 85 insertions(+), 97 deletions(-) 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 79a9972901..0f836947c1 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -6,9 +6,8 @@ import torch -from lightllm.common.basemodel import InferStateInfo, TransformerLayerInferTpl +from lightllm.common.basemodel import 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, @@ -21,7 +20,6 @@ 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 @@ -63,19 +61,6 @@ 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_, @@ -166,33 +151,11 @@ 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 = self._sparse_mla_bmm( - x=q, projection=layer_weight.k_b_proj_, infer_state=infer_state, out_features=self.kv_lora_rank - ).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"], @@ -313,11 +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 = self._sparse_mla_bmm( - x=input, projection=layer_weight.v_b_proj_, infer_state=infer_state, out_features=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 ) - output = layer_weight.o_weight_.mm(input.reshape(-1, self.tp_q_head_num_ * self.v_head_dim)) + # 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): diff --git a/unit_tests/models/glm5_next/test_layers.py b/unit_tests/models/glm5_next/test_layers.py index 8ac9c2d473..1aca878b2d 100644 --- a/unit_tests/models/glm5_next/test_layers.py +++ b/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] 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 index 5de243d770..fa9e727f87 100644 --- a/unit_tests/models/glm5_next/test_token_major_mla_bmm.py +++ b/unit_tests/models/glm5_next/test_token_major_mla_bmm.py @@ -11,59 +11,82 @@ 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]) +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)]) -def test_token_major_projection_matches_reference_and_keeps_guard_rows(rows: int, inner: int, cols: int) -> None: +@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, 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) + 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, 64, cols), 113.0, device="cuda", dtype=torch.bfloat16) + 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 = layer._sparse_mla_bmm(x=x, projection=projection, infer_state=infer, out_features=cols) - assert actual is output and actual.is_contiguous() + 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() - 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() + + +@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)