feat: add hadamard_transform kernel

Author: niushengxiao

lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py

@@ -251,7 +251,7 @@ def _get_indices(
     @staticmethod
     def _rotate_activation(x: torch.Tensor) -> torch.Tensor:
         assert x.dtype == torch.bfloat16
-        from sgl_kernel import hadamard_transform
+        from lightllm.models.deepseek3_2.triton_kernel.hadamard_transform import hadamard_transform
 
         hidden_size = x.size(-1)
         assert (hidden_size & (hidden_size - 1)) == 0, "Hidden size must be a power of 2 for Hadamard transform."

lightllm/models/deepseek3_2/triton_kernel/hadamard_transform.py

@@ -0,0 +1,57 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _butterfly_stage(x, GROUPS: tl.constexpr, STEP: tl.constexpr, BLOCK_N: tl.constexpr):
+    x_grouped = tl.reshape(x, (GROUPS, 2, STEP))
+    x_grouped = tl.permute(x_grouped, (0, 2, 1))
+    left, right = tl.split(x_grouped)
+    x_pair = tl.join(left + right, left - right)
+    x_pair = tl.permute(x_pair, (0, 2, 1))
+    return tl.reshape(x_pair, (BLOCK_N,))
+
+
+@triton.jit
+def _hadamard_transform_kernel(
+    X,
+    Y,
+    scale: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offsets = tl.arange(0, BLOCK_N)
+    x = tl.load(X + row * BLOCK_N + offsets).to(tl.float32)
+
+    x = _butterfly_stage(x, 64, 1, BLOCK_N)
+    x = _butterfly_stage(x, 32, 2, BLOCK_N)
+    x = _butterfly_stage(x, 16, 4, BLOCK_N)
+    x = _butterfly_stage(x, 8, 8, BLOCK_N)
+    x = _butterfly_stage(x, 4, 16, BLOCK_N)
+    x = _butterfly_stage(x, 2, 32, BLOCK_N)
+    x = _butterfly_stage(x, 1, 64, BLOCK_N)
+
+    tl.store(Y + row * BLOCK_N + offsets, x * scale)
+
+
+def hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
+    assert x.is_cuda, "hadamard_transform only supports CUDA tensors"
+    assert x.dtype == torch.bfloat16, "hadamard_transform expects bfloat16 input"
+
+    original_shape = x.shape
+    hidden_size = x.size(-1)
+    assert hidden_size == 128, "DeepSeek-V3.2 Hadamard transform expects hidden size 128"
+
+    x = x.contiguous()
+    out = torch.empty_like(x)
+    rows = x.numel() // hidden_size
+    _hadamard_transform_kernel[(rows,)](
+        x,
+        out,
+        scale,
+        BLOCK_N=hidden_size,
+        num_warps=4,
+    )
+
+    return out.view(original_shape)

lightllm/utils/backend_validator.py

@@ -196,12 +196,15 @@ def _validate_flashmla_sparse():
     except Exception as e:
         return False, f"sgl_kernel.flash_mla import failed: {type(e).__name__}: {e}"
 
-    batch, heads, seq, dim = 1, 64, 128, 512 + 64
+    batch, heads, seq = 1, 64, 128
+    kv_lora_rank = 512
+    qk_rope_head_dim = 64
+    qk_dim = kv_lora_rank + qk_rope_head_dim
     dtype = torch.bfloat16
     device = "cuda"
 
-    q = torch.randn(batch * seq, heads, dim, dtype=dtype, device=device)
-    kv = torch.zeros(batch * seq, 1, dim, dtype=dtype, device=device)
+    q = torch.randn(batch * seq, heads, qk_dim, dtype=dtype, device=device)
+    kv = torch.zeros(batch * seq, 1, qk_dim, dtype=dtype, device=device)
 
     index_topk = 128
     topk_indices = torch.zeros(batch * seq, index_topk, dtype=torch.int32, device=device)
@@ -210,8 +213,7 @@ def _validate_flashmla_sparse():
 
     topk_indices = topk_indices.view(batch * seq, 1, index_topk)
 
-    softmax_scale = 1.0 / (dim ** 0.5)
-    kv_lora_rank = dim
+    softmax_scale = 1.0 / (qk_dim ** 0.5)
 
     try:
         mla_out, _, _ = flash_mla_sparse_fwd(

unit_tests/models/deepseek3_2/triton_kernel/test_hadamard_transform.py

@@ -0,0 +1,92 @@
+import pytest
+import torch
+
+from lightllm.models.deepseek3_2.triton_kernel.hadamard_transform import hadamard_transform
+
+
+TP = 8
+INDEX_N_HEADS = 64
+INDEX_HEAD_DIM = 128
+TP_INDEX_N_HEADS = INDEX_N_HEADS // TP
+SCALE = INDEX_HEAD_DIM ** -0.5
+
+
+def _get_sgl_kernel_hadamard_transform():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required for hadamard_transform comparison")
+    try:
+        from sgl_kernel import hadamard_transform as sgl_hadamard_transform
+    except ImportError:
+        pytest.skip("sgl_kernel.hadamard_transform is not available")
+    return sgl_hadamard_transform
+
+
+def _bench(fn, x, warmup=30, iters=300):
+    for _ in range(warmup):
+        fn(x, scale=SCALE)
+    torch.cuda.synchronize()
+
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iters):
+        y = fn(x, scale=SCALE)
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) / iters, y
+
+
+@pytest.mark.parametrize("tokens", [1, 16, 128, 512, 1024, 2048, 4096, 8192, 16384])
+def test_hadamard_transform_matches_sgl_kernel_deepseek_v32_shapes(tokens):
+    sgl_hadamard_transform = _get_sgl_kernel_hadamard_transform()
+
+    q = torch.randn(tokens, TP_INDEX_N_HEADS, INDEX_HEAD_DIM, dtype=torch.bfloat16, device="cuda")
+    k = torch.randn(tokens, INDEX_HEAD_DIM, dtype=torch.bfloat16, device="cuda")
+
+    q_expected = sgl_hadamard_transform(q, scale=SCALE)
+    q_actual = hadamard_transform(q, scale=SCALE)
+    k_expected = sgl_hadamard_transform(k, scale=SCALE)
+    k_actual = hadamard_transform(k, scale=SCALE)
+    torch.cuda.synchronize()
+
+    assert torch.equal(q_actual, q_expected)
+    assert torch.equal(k_actual, k_expected)
+
+
+def test_hadamard_transform_perf_report_deepseek_v32_shapes():
+    sgl_hadamard_transform = _get_sgl_kernel_hadamard_transform()
+
+    print(
+        "\nDeepSeek-V3.2 per-rank shapes with tp=8:"
+        "\n  q: [tokens, 8, 128]"
+        "\n  k: [tokens, 128]"
+        "\n\ntokens | q_diff | k_diff | sgl_q ms | tri_q ms | sgl_k ms | tri_k ms | tri(q+k) ms | slowdown q+k"
+    )
+
+    for tokens in [1, 16, 128, 512, 1024, 2048, 4096, 8192, 16384]:
+        q = torch.randn(tokens, TP_INDEX_N_HEADS, INDEX_HEAD_DIM, dtype=torch.bfloat16, device="cuda")
+        k = torch.randn(tokens, INDEX_HEAD_DIM, dtype=torch.bfloat16, device="cuda")
+
+        q_expected = sgl_hadamard_transform(q, scale=SCALE)
+        q_actual = hadamard_transform(q, scale=SCALE)
+        k_expected = sgl_hadamard_transform(k, scale=SCALE)
+        k_actual = hadamard_transform(k, scale=SCALE)
+        torch.cuda.synchronize()
+
+        q_diff = (q_expected.float() - q_actual.float()).abs().max().item()
+        k_diff = (k_expected.float() - k_actual.float()).abs().max().item()
+        sgl_q_ms, _ = _bench(sgl_hadamard_transform, q)
+        tri_q_ms, _ = _bench(hadamard_transform, q)
+        sgl_k_ms, _ = _bench(sgl_hadamard_transform, k)
+        tri_k_ms, _ = _bench(hadamard_transform, k)
+        sgl_sum_ms = sgl_q_ms + sgl_k_ms
+        tri_sum_ms = tri_q_ms + tri_k_ms
+
+        print(
+            f"{tokens:6d} | {q_diff:6.1g} | {k_diff:6.1g} | "
+            f"{sgl_q_ms:8.4f} | {tri_q_ms:8.4f} | {sgl_k_ms:8.4f} | {tri_k_ms:8.4f} | "
+            f"{tri_sum_ms:11.4f} | {tri_sum_ms / sgl_sum_ms:10.2f}x"
+        )
+
+        assert q_diff == 0
+        assert k_diff == 0