66 lines
2.0 KiB
Python
66 lines
2.0 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
import torch
|
|
|
|
from sglang.jit_kernel.utils import cache_once, load_jit
|
|
|
|
if TYPE_CHECKING:
|
|
from tvm_ffi.module import Module
|
|
|
|
|
|
@cache_once
|
|
def _jit_concat_mla_k_module() -> Module:
|
|
return load_jit(
|
|
"concat_mla_k",
|
|
cuda_files=["elementwise/concat_mla.cuh"],
|
|
cuda_wrappers=[("concat_mla_k", "ConcatMlaKKernel::run")],
|
|
)
|
|
|
|
|
|
@cache_once
|
|
def _jit_concat_mla_absorb_q_module() -> Module:
|
|
return load_jit(
|
|
"concat_mla_absorb_q",
|
|
cuda_files=["elementwise/concat_mla.cuh"],
|
|
cuda_wrappers=[("concat_mla_absorb_q", "ConcatMlaAbsorbQKernel::run")],
|
|
)
|
|
|
|
|
|
def concat_mla_k(k: torch.Tensor, k_nope: torch.Tensor, k_rope: torch.Tensor) -> None:
|
|
"""
|
|
Concatenate k_nope and k_rope into k for MLA (Multi-head Latent Attention).
|
|
|
|
This kernel efficiently broadcasts k_rope across all heads while copying
|
|
k_nope values directly.
|
|
|
|
Args:
|
|
k: Output tensor of shape [num_tokens, num_heads=128, k_head_dim=192], dtype=bfloat16
|
|
k_nope: Input tensor of shape [num_tokens, num_heads=128, nope_head_dim=128], dtype=bfloat16
|
|
k_rope: Input tensor of shape [num_tokens, 1, rope_head_dim=64], dtype=bfloat16
|
|
"""
|
|
module = _jit_concat_mla_k_module()
|
|
module.concat_mla_k(k, k_nope, k_rope)
|
|
|
|
|
|
def concat_mla_absorb_q(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Concatenate tensors a and b for MLA absorbed Q computation.
|
|
|
|
Args:
|
|
a: Input tensor of shape [dim_0, dim_1, a_last_dim], dtype=bfloat16
|
|
b: Input tensor of shape [dim_0, dim_1, b_last_dim], dtype=bfloat16
|
|
|
|
Returns:
|
|
Output tensor of shape [dim_0, dim_1, a_last_dim + b_last_dim], dtype=bfloat16
|
|
"""
|
|
out = torch.empty(
|
|
(*a.shape[:-1], a.shape[-1] + b.shape[-1]),
|
|
dtype=a.dtype,
|
|
device=a.device,
|
|
)
|
|
module = _jit_concat_mla_absorb_q_module()
|
|
module.concat_mla_absorb_q(a, b, out)
|
|
return out
|