sglang 0.5.2rc2__py3-none-any.whl → 0.5.3rc2__py3-none-any.whl

sglang/srt/layers/attention/npu_ops/mla_preprocess.py

@@ -0,0 +1,393 @@
+import torch
+import torch.nn.functional as F
+
+from sglang.srt.utils import get_bool_env_var, is_npu
+
+_is_npu = is_npu()
+_ENABLE_MLA_PREPROCESS_FLAG = get_bool_env_var("SGLANG_NPU_USE_MLAPO")
+_NPU_FORMAT_NZ = 29
+
+
+def is_mla_preprocess_enabled() -> bool:
+    return _is_npu and _ENABLE_MLA_PREPROCESS_FLAG
+
+
+if is_mla_preprocess_enabled():
+    import sgl_kernel_npu
+    import torch_npu
+
+    torch.npu.config.allow_internal_format = True
+    torch.npu.set_compile_mode(jit_compile=False)
+
+
+def round_up(val: int, align: int) -> int:
+    if align == 0:
+        return 0
+    return -(val // -align) * align
+
+
+def transdata(nd_mat, block_size: tuple = (16, 16)):
+    r = round_up(nd_mat.shape[0], block_size[0])
+    c = round_up(nd_mat.shape[1], block_size[1])
+    r_pad = r - nd_mat.shape[0]
+    c_pad = c - nd_mat.shape[1]
+    nd_mat = F.pad(nd_mat, ((0, r_pad, 0, c_pad)))
+    nz_mat = torch.permute(
+        torch.reshape(
+            nd_mat,
+            (r // block_size[0], block_size[0], c // block_size[1], block_size[1]),
+        ),
+        [2, 0, 1, 3],
+    )
+    nz_mat = torch.reshape(
+        nz_mat, (nz_mat.shape[0], nz_mat.shape[1] * nz_mat.shape[2], nz_mat.shape[3])
+    )
+    return nz_mat
+
+
+def trans_rope_weight(weight, rope_dim):
+    weight_1 = weight[..., -rope_dim::2, :].contiguous()
+    weight_2 = weight[..., -rope_dim + 1 :: 2, :].contiguous()
+    weight[..., -rope_dim:, :] = torch.cat([weight_1, weight_2], dim=-2)
+
+    return weight.contiguous()
+
+
+class NPUFusedMLAPreprocess(torch.nn.Module):
+    def __init__(
+        self,
+        fused_qkv_a_proj_with_mqa,
+        q_a_layernorm,
+        kv_a_layernorm,
+        q_b_proj,
+        w_kc,
+        rotary_emb,
+        layer_id,
+        num_local_heads,
+        qk_nope_head_dim,
+        qk_rope_head_dim,
+    ):
+        super().__init__()
+        self.qkv_a_proj = fused_qkv_a_proj_with_mqa
+        self.q_a_layernorm = q_a_layernorm
+        self.kv_a_layernorm = kv_a_layernorm
+        self.q_b_proj = q_b_proj
+        self.w_kc = w_kc.contiguous()
+        self.rotary_emb = rotary_emb
+        self.layer_id = layer_id
+        self.has_preprocess_weights = False
+        self.dtype = None
+
+        self.q_lora_rank = self.q_b_proj.input_size  # 1536
+        self.kv_lora_rank = self.kv_a_layernorm.hidden_size  # 512
+        self.num_local_heads = num_local_heads  # tp
+        self.qk_nope_head_dim = qk_nope_head_dim  # 128
+        self.qk_rope_head_dim = qk_rope_head_dim  # 64
+        self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
+
+    def preprocess_weights(self, hidden_states):
+        self.dummy = torch.empty(
+            (hidden_states.shape[-1]),
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+        self.qkv_a_proj_input_offset = self.qkv_a_proj.input_offset.to(dtype=torch.int8)
+        self.q_b_proj_input_offset = self.q_b_proj.input_offset.to(dtype=torch.int8)
+
+        # matmul_0 weight [7168, 2112]
+        fused_qkv_a_proj_with_mqa_weight_q = self.qkv_a_proj.weight.data[
+            :, : self.q_lora_rank
+        ].clone()  # [7168, 1536]
+        fused_qkv_a_proj_with_mqa_weight_kv = self.qkv_a_proj.weight.data[
+            :, self.q_lora_rank :
+        ].clone()  # [7168, 576]
+        # rope fit
+        fused_qkv_a_proj_with_mqa_weight_kv_t = (
+            fused_qkv_a_proj_with_mqa_weight_kv.t().contiguous()
+        )
+        fused_qkv_a_proj_with_mqa_weight_kv_t = trans_rope_weight(
+            fused_qkv_a_proj_with_mqa_weight_kv_t, self.qk_rope_head_dim
+        )
+        fused_qkv_a_proj_with_mqa_weight_kv = (
+            fused_qkv_a_proj_with_mqa_weight_kv_t.t().contiguous()
+        )
+        # cat nz
+        fused_qkv_a_proj_with_mqa_weight_new = torch.cat(
+            (fused_qkv_a_proj_with_mqa_weight_kv, fused_qkv_a_proj_with_mqa_weight_q),
+            dim=-1,
+        )
+        fused_qkv_a_proj_with_mqa_weight = (
+            fused_qkv_a_proj_with_mqa_weight_new.t().contiguous()
+        )
+        fused_qkv_a_proj_with_mqa_weight_nz = (
+            transdata(fused_qkv_a_proj_with_mqa_weight, block_size=(16, 32))
+            .unsqueeze(0)
+            .contiguous()
+        )
+        self.qkv_a_proj_weight_nz = torch_npu.npu_format_cast(
+            fused_qkv_a_proj_with_mqa_weight_nz, _NPU_FORMAT_NZ
+        )
+
+        # matmul_0 deq_scale [2112]
+        fused_qkv_a_proj_with_mqa_deq_scale_q = self.qkv_a_proj.deq_scale.data[
+            : self.q_lora_rank
+        ].clone()  # [7168, 1536]
+        fused_qkv_a_proj_with_mqa_deq_scale_kv = self.qkv_a_proj.deq_scale.data[
+            self.q_lora_rank :
+        ].clone()  # [7168, 576]
+        # rope fit
+        fused_qkv_a_proj_with_mqa_deq_scale_kv = (
+            fused_qkv_a_proj_with_mqa_deq_scale_kv.reshape(
+                self.kv_lora_rank + self.qk_rope_head_dim, -1
+            ).contiguous()
+        )
+        fused_qkv_a_proj_with_mqa_deq_scale_kv = trans_rope_weight(
+            fused_qkv_a_proj_with_mqa_deq_scale_kv, self.qk_rope_head_dim
+        )
+        fused_qkv_a_proj_with_mqa_deq_scale_kv = (
+            fused_qkv_a_proj_with_mqa_deq_scale_kv.view(
+                self.kv_lora_rank + self.qk_rope_head_dim
+            ).contiguous()
+        )
+        self.qkv_a_proj_deq_scale_kvq = torch.cat(
+            (
+                fused_qkv_a_proj_with_mqa_deq_scale_kv,
+                fused_qkv_a_proj_with_mqa_deq_scale_q,
+            ),
+            dim=-1,
+        )
+
+        # matmul_0 quant_bias [2112]
+        fused_qkv_a_proj_with_mqa_quant_bias_q = self.qkv_a_proj.quant_bias.data[
+            : self.q_lora_rank
+        ].clone()  # [7168, 1536]
+        fused_qkv_a_proj_with_mqa_quant_bias_kv = self.qkv_a_proj.quant_bias.data[
+            self.q_lora_rank :
+        ].clone()  # [7168, 576]
+        # rope fit
+        fused_qkv_a_proj_with_mqa_quant_bias_kv = (
+            fused_qkv_a_proj_with_mqa_quant_bias_kv.reshape(
+                self.kv_lora_rank + self.qk_rope_head_dim, -1
+            ).contiguous()
+        )
+        fused_qkv_a_proj_with_mqa_quant_bias_kv = trans_rope_weight(
+            fused_qkv_a_proj_with_mqa_quant_bias_kv, self.qk_rope_head_dim
+        )
+        fused_qkv_a_proj_with_mqa_quant_bias_kv = (
+            fused_qkv_a_proj_with_mqa_quant_bias_kv.view(
+                self.kv_lora_rank + self.qk_rope_head_dim
+            ).contiguous()
+        )
+        self.qkv_a_proj_quant_bias_kvq = torch.cat(
+            (
+                fused_qkv_a_proj_with_mqa_quant_bias_kv,
+                fused_qkv_a_proj_with_mqa_quant_bias_q,
+            ),
+            dim=-1,
+        )
+
+        # matmul_1 weight [1536, num_head * 192]
+        q_b_proj_weight = self.q_b_proj.weight.data.clone()
+        q_b_proj_weight = q_b_proj_weight.t().reshape(
+            self.num_local_heads, self.qk_nope_head_dim + self.qk_rope_head_dim, -1
+        )
+        q_b_proj_weight = trans_rope_weight(q_b_proj_weight, self.qk_rope_head_dim)
+        q_b_proj_weight = q_b_proj_weight.reshape(
+            self.num_local_heads * (self.qk_nope_head_dim + self.qk_rope_head_dim), -1
+        )
+        q_b_proj_weight_nz = (
+            transdata(q_b_proj_weight, block_size=(16, 32)).unsqueeze(0).contiguous()
+        )
+        self.q_b_proj_weight_nz = torch_npu.npu_format_cast(
+            q_b_proj_weight_nz, _NPU_FORMAT_NZ
+        )
+
+        # matmul_1 deq_scale [num_head * 192]
+        q_b_proj_deq_scale = self.q_b_proj.deq_scale.data.clone()
+        q_b_proj_deq_scale = q_b_proj_deq_scale.reshape(
+            self.num_local_heads, self.qk_nope_head_dim + self.qk_rope_head_dim, -1
+        )
+        q_b_proj_deq_scale = trans_rope_weight(
+            q_b_proj_deq_scale, self.qk_rope_head_dim
+        )
+        self.q_b_proj_deq_scale = q_b_proj_deq_scale.reshape(
+            self.num_local_heads * (self.qk_nope_head_dim + self.qk_rope_head_dim)
+        )
+
+        # matmul_1 quant_bias [num_head * 192]
+        q_b_proj_quant_bias = self.q_b_proj.quant_bias.data.clone()
+        q_b_proj_quant_bias = q_b_proj_quant_bias.reshape(
+            self.num_local_heads, self.qk_nope_head_dim + self.qk_rope_head_dim, -1
+        )
+        q_b_proj_quant_bias = trans_rope_weight(
+            q_b_proj_quant_bias, self.qk_rope_head_dim
+        )
+        self.q_b_proj_quant_bias = q_b_proj_quant_bias.reshape(
+            self.num_local_heads * (self.qk_nope_head_dim + self.qk_rope_head_dim)
+        )
+
+    def get_sin_cos(self, positions):
+        cos_sin = self.rotary_emb.cos_sin_cache[positions]
+        cos, sin = cos_sin.chunk(2, dim=-1)
+        cos = cos.repeat(1, 2)
+        sin = sin.repeat(1, 2)
+        return cos, sin
+
+    def get_kv_cache_and_cache_idx(self, forward_batch):
+        k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer(self.layer_id)
+        slot_mapping = forward_batch.out_cache_loc.to(dtype=torch.int32)
+        return k_cache, v_cache, slot_mapping
+
+    def forward_absorb_prepare_npu_rms_norm_cache(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        forward_batch,
+        zero_allocator,
+    ):
+        bsz, _ = hidden_states.view(-1, hidden_states.shape[-1]).shape
+        self.dtype = hidden_states.dtype
+        self.cos, self.sin = self.get_sin_cos(positions)
+        self.kvCache, self.kvCacheRope, self.slotmapping = (
+            self.get_kv_cache_and_cache_idx(forward_batch)
+        )
+
+        if not self.has_preprocess_weights:
+            self.has_preprocess_weights = True
+
+        cos, sin = self.cos, self.sin
+
+        if self.q_lora_rank is not None:
+            fused_qkv_a_proj_out = self.qkv_a_proj(hidden_states)[0]
+            q_lowrank, latent_cache = fused_qkv_a_proj_out.split(
+                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1
+            )
+            q = self.q_a_layernorm(q_lowrank)
+            q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
+        else:
+            q = self.q_proj(hidden_states)[0].view(
+                -1, self.num_local_heads, self.qk_head_dim
+            )
+            latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
+
+        q_nope, q_pe = torch.split(
+            q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
+        )  # b*s,n,d
+
+        q_nope = q_nope.view(-1, self.num_local_heads, self.qk_nope_head_dim)
+        q_nope = torch.matmul(q_nope.transpose(0, 1), self.w_kc).transpose(0, 1)
+
+        q_pe = q_pe.view(-1, self.num_local_heads, 1, self.qk_rope_head_dim)
+        cos = cos.view(-1, 1, 1, self.qk_rope_head_dim)
+        sin = sin.view(-1, 1, 1, self.qk_rope_head_dim)
+        q_pe = torch_npu.npu_interleave_rope(q_pe, cos, sin)  # (B,N,S,D)
+        q_pe = q_pe.view(cos.shape[0], self.num_local_heads, self.qk_rope_head_dim)
+
+        latent_cache = latent_cache.view(
+            -1, 1, 1, self.kv_lora_rank + self.qk_rope_head_dim
+        )  # (B*S,N,1,D)
+
+        cache_mode = "PA_BNSD"
+        self.kvCache = self.kvCache.view(
+            -1,
+            forward_batch.attn_backend.page_size,
+            1,
+            forward_batch.attn_backend.kv_lora_rank,
+        )
+        self.kvCacheRope = self.kvCacheRope.view(
+            -1,
+            forward_batch.attn_backend.page_size,
+            1,
+            forward_batch.attn_backend.qk_rope_head_dim,
+        )
+        k_rope, k_nope, _, _ = torch_npu.npu_kv_rmsnorm_rope_cache(
+            latent_cache,
+            self.kv_a_layernorm.weight,
+            cos,
+            sin,
+            self.slotmapping.to(torch.int64),
+            self.kvCacheRope,
+            self.kvCache,
+            epsilon=self.kv_a_layernorm.variance_epsilon,
+            cache_mode=cache_mode,
+        )
+
+        return (q_pe, k_rope, q_nope, k_nope, forward_batch, zero_allocator, positions)
+
+    def forward_mlapo(self, positions, hidden_states, forward_batch, zero_allocator):
+        input_dtype = hidden_states.dtype
+        if not self.has_preprocess_weights:
+            self.preprocess_weights(hidden_states)
+            self.has_preprocess_weights = True
+            self.dtype = hidden_states.dtype
+
+        cos, sin = self.get_sin_cos(positions)
+        k_cache, v_cache, slot_mapping = self.get_kv_cache_and_cache_idx(forward_batch)
+
+        q_nope_out = torch.empty(
+            (hidden_states.shape[0], self.w_kc.shape[0], k_cache.shape[-1]),
+            dtype=input_dtype,
+            device=hidden_states.device,
+        )
+        q_rope_out = torch.empty(
+            (hidden_states.shape[0], self.w_kc.shape[0], v_cache.shape[-1]),
+            dtype=input_dtype,
+            device=hidden_states.device,
+        )
+
+        # TODO: dummy inputs to be removed
+        # https://github.com/sgl-project/sgl-kernel-npu/issues/78
+        torch.ops.npu.mla_preprocess(
+            hidden_states,
+            self.dummy,
+            self.dummy,
+            self.qkv_a_proj_weight_nz,
+            self.qkv_a_proj_deq_scale_kvq,
+            self.q_a_layernorm.weight,
+            self.q_a_layernorm.bias,
+            self.q_b_proj_weight_nz,
+            self.q_b_proj_deq_scale,
+            self.kv_a_layernorm.weight,
+            cos,
+            sin,
+            self.w_kc,
+            k_cache,
+            v_cache,
+            slot_mapping,
+            quant_scale0=self.qkv_a_proj.input_scale,
+            quant_offset0=self.qkv_a_proj_input_offset,
+            bias0=self.qkv_a_proj_quant_bias_kvq,
+            quant_scale1=self.q_b_proj.input_scale,
+            quant_offset1=self.q_b_proj_input_offset,
+            bias1=self.q_b_proj_quant_bias,
+            cache_mode="krope_ctkv",
+            quant_mode="per_tensor_quant_asymm",
+            q_out0=q_nope_out,
+            kv_cache_out0=k_cache,
+            q_out1=q_rope_out,
+            kv_cache_out1=v_cache,
+        )
+        return (
+            q_rope_out,
+            v_cache,
+            q_nope_out,
+            k_cache,
+            forward_batch,
+            zero_allocator,
+            positions,
+        )
+
+    def forward(self, positions, hidden_states, forward_batch, zero_allocator):
+        _is_w8a8 = (
+            hasattr(self.qkv_a_proj.quant_method, "quantization_config")
+            and self.qkv_a_proj.quant_method.quantization_config.get_name()
+            == "w8a8_int8"
+        )
+        if _is_w8a8:
+            return self.forward_mlapo(
+                positions, hidden_states, forward_batch, zero_allocator
+            )
+        else:
+            return self.forward_absorb_prepare_npu_rms_norm_cache(
+                positions, hidden_states, forward_batch, zero_allocator
+            )

sglang/srt/layers/attention/nsa/dequant_k_cache.py

@@ -0,0 +1,163 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.layers.attention.nsa.utils import NSA_DEQUANT_K_CACHE_FAST
+
+
+def dequantize_k_cache(quant_k_cache):
+    if NSA_DEQUANT_K_CACHE_FAST:
+        return _dequantize_k_cache_fast_wrapped(quant_k_cache)
+    else:
+        return _dequantize_k_cache_slow(quant_k_cache)
+
+
+def _dequantize_k_cache_slow(
+    quant_k_cache: torch.Tensor,  # (num_blocks, block_size, 1, bytes_per_token)
+    dv: int = 512,
+    tile_size: int = 128,
+    d: int = 576,
+) -> torch.Tensor:
+    """
+    De-quantize the k-cache
+    """
+    assert dv % tile_size == 0
+    num_tiles = dv // tile_size
+    num_blocks, block_size, h_k, _ = quant_k_cache.shape
+    assert h_k == 1
+    result = torch.empty(
+        (num_blocks, block_size, d), dtype=torch.bfloat16, device=quant_k_cache.device
+    )
+
+    quant_k_cache = quant_k_cache.view(num_blocks, block_size, -1)
+
+    input_nope = quant_k_cache[..., :dv]
+    input_scale = quant_k_cache[..., dv : dv + num_tiles * 4].view(torch.float32)
+    input_rope = quant_k_cache[..., dv + num_tiles * 4 :].view(torch.bfloat16)
+    result[..., dv:] = input_rope
+
+    for tile_idx in range(0, num_tiles):
+        cur_nope = input_nope[
+            ..., tile_idx * tile_size : (tile_idx + 1) * tile_size
+        ].to(torch.float32)
+        cur_scales = input_scale[..., tile_idx].unsqueeze(-1)
+        result[..., tile_idx * tile_size : (tile_idx + 1) * tile_size] = (
+            cur_nope * cur_scales
+        )
+
+    result = result.view(num_blocks, block_size, 1, d)
+    return result
+
+
+def _dequantize_k_cache_fast_wrapped(
+    quant_k_cache: torch.Tensor,
+    dv: int = 512,
+    tile_size: int = 128,
+) -> torch.Tensor:
+    # TODO the final API may be 2D instead of 4D, thus we convert them here
+    num_blocks, block_size, _, dim_quant = quant_k_cache.shape
+    assert dv == 512
+    assert dim_quant == 656
+    assert tile_size == 128
+    quant_k_cache = quant_k_cache.view((-1, dim_quant))
+
+    output = _dequantize_k_cache_fast(quant_k_cache)
+
+    return output.view(num_blocks, block_size, 1, -1)
+
+
+def _dequantize_k_cache_fast(quant_k_cache, group_size: int = 128):
+    num_tokens, dim_quant = quant_k_cache.shape
+
+    assert quant_k_cache.dtype == torch.float8_e4m3fn
+    dim_nope = 512
+    dim_rope = 64
+    num_tiles = dim_nope // group_size
+    assert dim_quant == 656
+
+    output = torch.empty(
+        (num_tokens, dim_nope + dim_rope),
+        dtype=torch.bfloat16,
+        device=quant_k_cache.device,
+    )
+
+    num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
+    assert num_blocks_per_token == 5
+
+    assert dim_nope % group_size == 0
+    NUM_NOPE_BLOCKS = dim_nope // group_size
+
+    input_nope_q = quant_k_cache[:, :dim_nope]
+    input_nope_s = quant_k_cache[:, dim_nope : dim_nope + num_tiles * 4].view(
+        torch.float32
+    )
+    input_rope = quant_k_cache[:, dim_nope + num_tiles * 4 :].view(torch.bfloat16)
+
+    _dequantize_k_cache_fast_kernel[(num_tokens, num_blocks_per_token)](
+        output,
+        input_nope_q,
+        input_nope_s,
+        input_rope,
+        output.stride(0),
+        input_nope_q.stride(0),
+        input_nope_s.stride(0),
+        input_rope.stride(0),
+        NUM_NOPE_BLOCKS=NUM_NOPE_BLOCKS,
+        GROUP_SIZE=group_size,
+        DIM_NOPE=dim_nope,
+        DIM_ROPE=dim_rope,
+    )
+
+    return output
+
+
+@triton.jit
+def _dequantize_k_cache_fast_kernel(
+    output_ptr,
+    input_nope_q_ptr,
+    input_nope_s_ptr,
+    input_rope_ptr,
+    output_stride_0: int,
+    input_nope_q_stride_0: int,
+    input_nope_s_stride_0: int,
+    input_rope_stride_0: int,
+    NUM_NOPE_BLOCKS: tl.constexpr,
+    GROUP_SIZE: tl.constexpr,
+    DIM_NOPE: tl.constexpr,
+    DIM_ROPE: tl.constexpr,
+):
+    token_id = tl.program_id(0)
+    raw_block_id = tl.program_id(1)
+
+    if raw_block_id < NUM_NOPE_BLOCKS:
+        # a. dequant nope
+        effective_block_id = raw_block_id
+
+        offs_q = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
+        mask = offs_q < DIM_NOPE
+        ptr_q = input_nope_q_ptr + token_id * input_nope_q_stride_0 + offs_q
+        ptr_s = input_nope_s_ptr + token_id * input_nope_s_stride_0 + effective_block_id
+
+        y_q = tl.load(ptr_q, mask=mask, other=0.0).to(tl.float32)
+        y_s = tl.load(ptr_s)
+
+        y = (y_q * y_s).to(output_ptr.dtype.element_ty)
+
+        dst_ptr = output_ptr + token_id * output_stride_0 + offs_q
+        tl.store(dst_ptr, y, mask=mask)
+    else:
+        # b. copy rope
+        effective_block_id = raw_block_id - NUM_NOPE_BLOCKS
+
+        offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
+        mask = offs < DIM_ROPE
+
+        src_ptr = input_rope_ptr + token_id * input_rope_stride_0 + offs
+        dst_ptr = output_ptr + token_id * output_stride_0 + DIM_NOPE + offs
+
+        data = tl.load(src_ptr, mask=mask).to(tl.bfloat16)
+        tl.store(dst_ptr, data, mask=mask)
+
+
+if __name__ == "__main__":
+    raise Exception("UT is in quant_k_cache.py")