From 87deeb05a42ace79dccaef84075c8a9aa31b2040 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Wed, 11 Nov 2026 15:19:59 +0000 Subject: [PATCH 1/2] perf(kda): parallelize diagonal solves in Prefill Solve the four diagonal blocks independently and retain FP32 inverse tiles through merging. Select by chunk occupancy to bound scratch, with CUDA regression coverage for packed layouts, dtypes, precision modes and fallback behavior. --- .../triton_kernel/linear_att/fla/ops/kda.py | 1 + .../linear_att/fla/ops/solve_tril.py | 125 ++++++++++++------ .../linear_att/test_kda_autotune.py | 70 ++++++++++ 3 files changed, 156 insertions(+), 40 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py index 6d627078dd..8555b83952 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py @@ -1122,6 +1122,7 @@ def _chunk_kda_fwd_with_cumulative_g( A=Akk, cu_seqlens=cu_seqlens, output_dtype=k.dtype, + split_diagonal=True, ) del Akk # 3. Prepare U/W/Kg for the state update without reading the input state S_in. diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py index b5b6cfc369..c061d3ef25 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py @@ -199,6 +199,27 @@ def merge_16x16_to_32x32_inverse_kernel( desc_o.store([i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne")) +@triton.jit +def _invert_64_diagonal_tiles(A, Diag, cu_seqlens, chunk_indices, H: tl.constexpr): + """Invert the four independent 16x16 blocks before the full inverse merge.""" + i_t, h, d = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + chunk = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos = tl.load(cu_seqlens + i_n).to(tl.int32) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + row = chunk * 64 + d * 16 + i = tl.arange(0, 16) + base = A + ((bos + row).to(tl.int64) * H + h) * 64 + d * 16 + tile = tl.load(base + i[:, None] * H * 64 + i[None, :], mask=(bos + row + i[:, None]) < eos, other=0).to(tl.float32) + inv = -tl.where(i[:, None] > i[None, :], tile, 0) + for j in range(2, min(16, eos - bos - row)): + a = -tl.load(base + j * H * 64 + i) + a += tl.sum(a[:, None] * inv, 0) + inv = tl.where((i == j)[:, None], a, inv) + inv += i[:, None] == i[None, :] + tl.store(Diag + ((i_t.to(tl.int64) * H + h) * 4 + d) * 256 + i[:, None] * 16 + i[None, :], inv) + + @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @triton.jit(do_not_specialize=["T"]) def merge_16x16_to_64x64_inverse_kernel( @@ -212,6 +233,8 @@ def merge_16x16_to_64x64_inverse_kernel( USE_TMA: tl.constexpr, IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, + Diag=None, + PREINVERTED: tl.constexpr = False, ): i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H @@ -234,49 +257,57 @@ def merge_16x16_to_64x64_inverse_kernel( A += (bos * H + i_h) * BT Ai += (bos * H + i_h) * BT - if not USE_TMA: - p_A_11 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) - p_A_22 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) - p_A_33 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0)) - p_A_44 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0)) - b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) - b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) - b_Ai_33 = tl.load(p_A_33, boundary_check=(0, 1)).to(tl.float32) - b_Ai_44 = tl.load(p_A_44, boundary_check=(0, 1)).to(tl.float32) - else: + if USE_TMA: desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16]) - b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32) - b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32) - b_Ai_33 = desc.load([i_t * BT + 32, 32]).to(tl.float32) - b_Ai_44 = desc.load([i_t * BT + 48, 48]).to(tl.float32) - - # [16, 16] - b_Ai_11 = -tl.where(m_A, b_Ai_11, 0) - b_Ai_22 = -tl.where(m_A, b_Ai_22, 0) - b_Ai_33 = -tl.where(m_A, b_Ai_33, 0) - b_Ai_44 = -tl.where(m_A, b_Ai_44, 0) + if PREINVERTED: + base = Diag + (tl.program_id(0).to(tl.int64) * H + i_bh) * 1024 + o_i[:, None] * 16 + o_i[None, :] + b_Ai_11 = tl.load(base) + b_Ai_22 = tl.load(base + 256) + b_Ai_33 = tl.load(base + 512) + b_Ai_44 = tl.load(base + 768) + else: + if not USE_TMA: + p_A_11 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) + p_A_22 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) + p_A_33 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0)) + p_A_44 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0)) + b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) + b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) + b_Ai_33 = tl.load(p_A_33, boundary_check=(0, 1)).to(tl.float32) + b_Ai_44 = tl.load(p_A_44, boundary_check=(0, 1)).to(tl.float32) + else: + b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32) + b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32) + b_Ai_33 = desc.load([i_t * BT + 32, 32]).to(tl.float32) + b_Ai_44 = desc.load([i_t * BT + 48, 48]).to(tl.float32) - for i in range(2, min(16, T - i_t * BT)): - b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i) - b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0) - b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11) - for i in range(16 + 2, min(32, T - i_t * BT)): - b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16) - b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0) - b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22) - for i in range(32 + 2, min(48, T - i_t * BT)): - b_a_33 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 32) - b_a_33 += tl.sum(b_a_33[:, None] * b_Ai_33, 0) - b_Ai_33 = tl.where((o_i == i - 32)[:, None], b_a_33, b_Ai_33) - for i in range(48 + 2, min(64, T - i_t * BT)): - b_a_44 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 48) - b_a_44 += tl.sum(b_a_44[:, None] * b_Ai_44, 0) - b_Ai_44 = tl.where((o_i == i - 48)[:, None], b_a_44, b_Ai_44) - b_Ai_11 += m_I - b_Ai_22 += m_I - b_Ai_33 += m_I - b_Ai_44 += m_I + # [16, 16] + b_Ai_11 = -tl.where(m_A, b_Ai_11, 0) + b_Ai_22 = -tl.where(m_A, b_Ai_22, 0) + b_Ai_33 = -tl.where(m_A, b_Ai_33, 0) + b_Ai_44 = -tl.where(m_A, b_Ai_44, 0) + + for i in range(2, min(16, T - i_t * BT)): + b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i) + b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0) + b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11) + for i in range(16 + 2, min(32, T - i_t * BT)): + b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16) + b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0) + b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22) + for i in range(32 + 2, min(48, T - i_t * BT)): + b_a_33 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 32) + b_a_33 += tl.sum(b_a_33[:, None] * b_Ai_33, 0) + b_Ai_33 = tl.where((o_i == i - 32)[:, None], b_a_33, b_Ai_33) + for i in range(48 + 2, min(64, T - i_t * BT)): + b_a_44 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 48) + b_a_44 += tl.sum(b_a_44[:, None] * b_Ai_44, 0) + b_Ai_44 = tl.where((o_i == i - 48)[:, None], b_a_44, b_Ai_44) + b_Ai_11 += m_I + b_Ai_22 += m_I + b_Ai_33 += m_I + b_Ai_44 += m_I if not USE_TMA: p_A_21 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) @@ -412,6 +443,8 @@ def solve_tril( A: torch.Tensor, cu_seqlens: torch.Tensor | None = None, output_dtype: torch.dtype = torch.float, + *, + split_diagonal: bool = False, ) -> torch.Tensor: """ Compute the inverse of the matrix I + A @@ -425,6 +458,9 @@ def solve_tril( output_dtype (torch.dtype): The dtype of the output tensor. Default: `torch.float`. If `None`, the output dtype will be the same as the input dtype. + split_diagonal (bool): + Split the four diagonal solves for packed 64-token chunks with + ``cu_seqlens`` and enough work to offset the extra kernel launch. Returns: (I + A)^-1 with the same shape as A @@ -448,6 +484,14 @@ def solve_tril( if is_tma_supported: _ensure_triton_allocator() + split_args = {} + # Bound diagonal scratch by the FP32 input size and skip sub-chunk batches. + if split_diagonal and (B, BT) == (1, 64) and T >= max(BT, 16 * NT): + # Keep diagonal inverses in FP32 until the final output conversion. + diag = torch.empty((NT, H, 4, 16, 16), dtype=torch.float32, device=A.device) + _invert_64_diagonal_tiles[NT, H, 4](A, diag, cu_seqlens, chunk_indices, H=H, num_warps=4) + split_args = {"Diag": diag, "PREINVERTED": True} + merge_fn[NT, B * H]( A=A, Ai=Ai, @@ -458,5 +502,6 @@ def solve_tril( BT=BT, USE_TMA=is_tma_supported, DOT_PRECISION=FLA_TRIL_PRECISION, + **split_args, ) return Ai diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py index 525dfae756..798c3a2c90 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py @@ -1,5 +1,6 @@ import collections import functools +import importlib import inspect import json from itertools import accumulate @@ -194,3 +195,72 @@ def test_kda_requires_packed_inputs(fused_gate, invalid_input): expected_message = "packed q" if invalid_input == "batch-dimension" else "cu_seqlens is required" with pytest.raises(AssertionError, match=expected_message): kernel(**inputs) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + "heads,input_dtype,output_dtype,precision", + [ + (2, torch.float32, torch.bfloat16, "ieee"), + (32, torch.bfloat16, torch.bfloat16, "ieee"), + (64, torch.float16, torch.float16, "ieee"), + (64, torch.float32, torch.float32, "ieee"), + (64, torch.float32, torch.bfloat16, "tf32"), + (64, torch.float32, torch.bfloat16, "tf32x3"), + ], +) +def test_split_diagonal_matches_unsplit(monkeypatch, heads, input_dtype, output_dtype, precision): + if precision != "ieee" and torch.cuda.get_device_capability()[0] < 8: + pytest.skip("TF32 requires SM80+") + solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril") + monkeypatch.setattr(solver, "FLA_TRIL_PRECISION", precision) + lengths = (1, 16, 1519) + cu_seqlens = torch.tensor([0, *accumulate(lengths)], device="cuda", dtype=torch.int32) + rows = torch.cat([torch.arange(length, device="cuda") % 64 for length in lengths]) + matrix = torch.randn(1, sum(lengths), heads, 64, device="cuda", dtype=input_dtype) * 0.025 + matrix.mul_((rows[:, None] > torch.arange(64, device="cuda"))[None, :, None, :]) + expected = solver.solve_tril(matrix, cu_seqlens, output_dtype) + launches = [] + original = solver._invert_64_diagonal_tiles.run + + def observe(*args, **kwargs): + launches.append(True) + return original(*args, **kwargs) + + monkeypatch.setattr(solver._invert_64_diagonal_tiles, "run", observe) + actual = solver.solve_tril(matrix, cu_seqlens, output_dtype, split_diagonal=True) + assert len(launches) == 1 + assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + "lengths,expect_split", + [ + ((512,) * 16, True), + ((64,) * 128, True), + ((16,) * 16, True), + ((128,) * 128, True), + ((1,) * 256, False), + ((32,), False), + ], + ids=["16-requests", "128-requests", "quarter-full-chunks", "16k-tokens", "mostly-empty-chunks", "tiny-batch"], +) +def test_split_diagonal_uses_chunk_occupancy(monkeypatch, lengths, expect_split): + solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril") + cu_seqlens = torch.tensor([0, *accumulate(lengths)], device="cuda", dtype=torch.int32) + rows = torch.cat([torch.arange(length, device="cuda") % 64 for length in lengths]) + matrix = torch.randn(1, sum(lengths), 2, 64, device="cuda") * 0.025 + matrix.mul_((rows[:, None] > torch.arange(64, device="cuda"))[None, :, None, :]) + expected = solver.solve_tril(matrix, cu_seqlens, torch.bfloat16) + launches = [] + original = solver._invert_64_diagonal_tiles.run + + def observe(*args, **kwargs): + launches.append(True) + return original(*args, **kwargs) + + monkeypatch.setattr(solver._invert_64_diagonal_tiles, "run", observe) + actual = solver.solve_tril(matrix, cu_seqlens, torch.bfloat16, split_diagonal=True) + assert len(launches) == int(expect_split) + assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) From f98c19353e6cdd6f1b66f3df39d9b3945d7241b4 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Thu, 12 Nov 2026 19:41:25 +0800 Subject: [PATCH 2/2] Simplify diagonal solver dispatch and support dense batches --- .../linear_att/fla/ops/solve_tril.py | 54 +++++++++++-------- .../linear_att/test_kda_autotune.py | 21 ++++++++ 2 files changed, 52 insertions(+), 23 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py index c061d3ef25..3512788934 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py @@ -199,14 +199,20 @@ def merge_16x16_to_32x32_inverse_kernel( desc_o.store([i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne")) -@triton.jit -def _invert_64_diagonal_tiles(A, Diag, cu_seqlens, chunk_indices, H: tl.constexpr): +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.jit(do_not_specialize=["T"]) +def _invert_64_diagonal_tiles(A, Diag, cu_seqlens, chunk_indices, T, H: tl.constexpr, IS_VARLEN: tl.constexpr): """Invert the four independent 16x16 blocks before the full inverse merge.""" - i_t, h, d = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) - chunk = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos = tl.load(cu_seqlens + i_n).to(tl.int32) - eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_t, i_bh, d = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + chunk = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos = tl.load(cu_seqlens + i_n).to(tl.int32) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + else: + bos, eos = i_b * T, (i_b + 1) * T + chunk = i_t row = chunk * 64 + d * 16 i = tl.arange(0, 16) base = A + ((bos + row).to(tl.int64) * H + h) * 64 + d * 16 @@ -217,7 +223,7 @@ def _invert_64_diagonal_tiles(A, Diag, cu_seqlens, chunk_indices, H: tl.constexp a += tl.sum(a[:, None] * inv, 0) inv = tl.where((i == j)[:, None], a, inv) inv += i[:, None] == i[None, :] - tl.store(Diag + ((i_t.to(tl.int64) * H + h) * 4 + d) * 256 + i[:, None] * 16 + i[None, :], inv) + tl.store(Diag + ((i_t.to(tl.int64) * tl.num_programs(1) + i_bh) * 4 + d) * 256 + i[:, None] * 16 + i[None, :], inv) @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @@ -261,7 +267,9 @@ def merge_16x16_to_64x64_inverse_kernel( desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16]) if PREINVERTED: - base = Diag + (tl.program_id(0).to(tl.int64) * H + i_bh) * 1024 + o_i[:, None] * 16 + o_i[None, :] + base = ( + Diag + (tl.program_id(0).to(tl.int64) * tl.num_programs(1) + i_bh) * 1024 + o_i[:, None] * 16 + o_i[None, :] + ) b_Ai_11 = tl.load(base) b_Ai_22 = tl.load(base + 256) b_Ai_33 = tl.load(base + 512) @@ -459,8 +467,8 @@ def solve_tril( The dtype of the output tensor. Default: `torch.float`. If `None`, the output dtype will be the same as the input dtype. split_diagonal (bool): - Split the four diagonal solves for packed 64-token chunks with - ``cu_seqlens`` and enough work to offset the extra kernel launch. + Split the four diagonal solves for 64-token chunks. Small inputs + and batches with excessive FP32 workspace use the original solver. Returns: (I + A)^-1 with the same shape as A @@ -473,24 +481,24 @@ def solve_tril( NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) Ai = torch.zeros_like(A, dtype=output_dtype) + # Ensure Triton allocator is set for TMA kernels that require scratch memory + if is_tma_supported: + _ensure_triton_allocator() + + split_args = {} if BT == 16: merge_fn = solve_tril_16x16_kernel elif BT == 32: merge_fn = merge_16x16_to_32x32_inverse_kernel elif BT == 64: merge_fn = merge_16x16_to_64x64_inverse_kernel - - # Ensure Triton allocator is set for TMA kernels that require scratch memory - if is_tma_supported: - _ensure_triton_allocator() - - split_args = {} - # Bound diagonal scratch by the FP32 input size and skip sub-chunk batches. - if split_diagonal and (B, BT) == (1, 64) and T >= max(BT, 16 * NT): - # Keep diagonal inverses in FP32 until the final output conversion. - diag = torch.empty((NT, H, 4, 16, 16), dtype=torch.float32, device=A.device) - _invert_64_diagonal_tiles[NT, H, 4](A, diag, cu_seqlens, chunk_indices, H=H, num_warps=4) - split_args = {"Diag": diag, "PREINVERTED": True} + # Bound the FP32 workspace by the FP32 input size and skip sub-chunk inputs. + diag_numel = NT * B * H * 4 * 16 * 16 + if split_diagonal and T >= BT and diag_numel <= A.numel(): + # Keep diagonal inverses in FP32 until the final output conversion. + diag = torch.empty((NT, B * H, 4, 16, 16), dtype=torch.float32, device=A.device) + _invert_64_diagonal_tiles[NT, B * H, 4](A, diag, cu_seqlens, chunk_indices, T, H=H, num_warps=4) + split_args = {"Diag": diag, "PREINVERTED": True} merge_fn[NT, B * H]( A=A, diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py index 798c3a2c90..8d92d6cc8d 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py @@ -233,6 +233,27 @@ def observe(*args, **kwargs): assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("batch,tokens,heads", [(1, 64, 2), (2, 65, 3), (3, 129, 2)]) +def test_split_diagonal_matches_unsplit_dense(monkeypatch, batch, tokens, heads): + solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril") + rows = torch.arange(tokens, device="cuda") % 64 + matrix = torch.randn(batch, tokens, heads, 64, device="cuda") * 0.025 + matrix.mul_((rows[:, None] > torch.arange(64, device="cuda"))[None, :, None, :]) + expected = solver.solve_tril(matrix, output_dtype=torch.bfloat16) + launches = [] + original = solver._invert_64_diagonal_tiles.run + + def observe(*args, **kwargs): + launches.append(True) + return original(*args, **kwargs) + + monkeypatch.setattr(solver._invert_64_diagonal_tiles, "run", observe) + actual = solver.solve_tril(matrix, output_dtype=torch.bfloat16, split_diagonal=True) + assert len(launches) == 1 + assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize( "lengths,expect_split",