Click the blue text to follow us
When adapting hybrid models on the inference framework (vLLM/SGLang), the key is to construct the corresponding linear modules. The linear module used by the QwenNext model is GDN (Gated Delta Networks), while KimiLinear uses KDA (Kimi Delta Attention). The basic structure of both modules is similar, both being improvements of linear attention based on the Delta rule. This article mainly focuses on GDN for explanation.
1 Construction of Formulas
In the inference of GDN and KDA, two types of formula forms are used:
Recursive and Chunkwise formulas. For the prefill of inference, when the input sequence is relatively short, recursive formulas are used, while for longer sequences, chunkwise methods are more efficient; during the decode phase, recursive formulas are generally sufficient. Below, we will analyze the calculation construction using GDN as an example.
1.1 Recursive Formulas
The calculation formula for updating the state S in GDN:
The calculation formula for the output:
To obtain the state value and output result, the input value must first be processed.
Mapping Calculation Q/K/V, G (supporting multiple heads, each head calculates the same) calculation formula:
Where the kernel function ф uses Silu, and Conv is a short convolution operation. Here, G is the gate controlling the output (not gamma decay).
# Convolution operation used in formula 4:conv1d = nn.Conv1d( in_channels=self.conv_dim, out_channels=self.conv_dim, bias=False, kernel_size=self.conv_kernel_size, groups=self.conv_dim, padding=self.conv_kernel_size - 1,)# Processing for formula 5:def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6): inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps) return x * inv_normCalculation formula for alpha:
Where A and Δ (dt_bias) are learnable parameters, and during the inference phase, values are directly loaded from the weights.
# Pseudo codeA = torch.empty(num_v_heads)A_log = nn.Parameter(torch.log(A))dt_bias = nn.Parameter(torch.ones(num_v_heads))# A_log, dt_bias need to be loaded from weights# Calculation of gamma value (formula 7):g = -A_log.float().exp() * F.softplus(a.float() + dt_bias)Calculation formula for beta:Formula 1’s S is split into two steps:
# Pseudo code# Iterative calculation process:out = torch.zeros(batch_size, num_heads, sequence_length, v_head_dim)state = torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim)for i in range(sequence_length): q_t = query[:, :, i] k_t = key[:, :, i] v_t = value[:, :, i] g_t = g[:, :, i].exp().unsqueeze(-1).unsqueeze(-1) beta_t = beta[:, :, i].unsqueeze(-1) kv_mem = (state * k_t.unsqueeze(-1)).sum(dim=-2) # Formula 10 v_new_t = v_t - kv_mem # Formula 9 state = state * g_t + k_t.unsqueeze(-1) * (v_new_t * beta_t).unsqueeze(-2) # Formula 2 out[:, :, i] = (state * q_t.unsqueeze(-1)).sum(dim=-2)1.2 Chunkwise Formulas
UT Transformation Formula:
Where M represents the mask operation, М indicates strict lower triangular operation (StrictTril), which can be directly called in torch using tril(-1). The gamma decay matrix is defined to satisfy:
# Pseudo code K_beta = K * beta.unsqueeze(-1) V_beta = V * beta.unsqueeze(-1) # Chunk decay calculation: g = g.cumsum(dim=-1) gamma = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril()).exp() # Formula 11: T = -(K_beta @ K.t() * gamma).tril(-1) for i in range(1, C): T[i, :i] = T[i, :i] + (T[i, :, None] * T[:, :i]).sum(-2) T += torch.eye(C) # Formula 12: W = T @ (K_beta * g.exp().unsqueeze(-1)) # Formula 13: U = T @ V_betaTo obtain γ from the input, the value passed to the function is generally the α value that has not undergone the exponential operation. The calculation can first perform accumulation and then calculate the exponential. It is evident that the former calculation is more efficient.
-
g.cumsum(dim=-1).exp() is equivalent to g.exp().cumprod();
-
(g1 – g2).exp() is equivalent to g1.exp() / g2.exp();
State Refresh and Result Output:
# Pseudo codedef chunk_gated_delta_rule_forward(Q, K, V, beta, g, C): """ Q/K/V: query, key, value of shape [L, d] beta: beta of shape [L] g: gate of shape [L], with exp() C: chunk size """ # L: sequence length, d: head dimension L, d = Q.shape # Chunking: Q, K, V = map(lambda x: x.reshape(-1, C, d), [Q, K, V]) beta = beta.reshape(-1, C) g = g.reshape(-1, C) K_beta = K * beta.unsqueeze(-1) V_beta = V * beta.unsqueeze(-1) # Chunk decay calculation: g = g.cumsum(dim=-1) gamma = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril()).exp() # Formula 11: T = -(K_beta @ K.t() * gamma).tril(-1) for i in range(1, C): T[i, :i] = T[i, :i] + (T[i, :, None] * T[:, :i]).sum(-2) T += torch.eye(C) # Formula 12: W = T @ (K_beta * g.exp().unsqueeze(-1)) # Formula 13: U = T @ V_beta S = torch.zeros(d, d) O = torch.zeros_like(U) for i in range(L//C): q_i, k_i, u_i, w_i = Q[i], K[i], U[i], W[i] # Formula 15: new_v_i = u_i - w_i @ S o_inter = (q_i * g[i].exp()) @ S o_intra = (q_i @ k_i.t() * gamma[:, i]).tril() # Formula 16: S = S * g[i, :].unsqueeze(-1) \ + (k_i * (g[i].unsqueeze(-1) - g[:, i]).exp()).t() @ new_v_i # Formula 17: O[i] = o_inter + o_intra @ new_v_i return O.reshape(L, d)2 Implementation of Code in the Framework
Current Status: The adaptation of GDN and KDA in vLLM/SGLang follows the existing logic of the Mamba model, with the main work being slight adjustments to cache management and the implementation of corresponding accelerated operators. In terms of functionality, the framework already has relevant code implementations; however, performance-wise, the open-source code is still undergoing optimization iterations.
Below, we will introduce KDA in SGLang as an example.
2.1 Operator Implementation
Multiple steps in the linear attention module can use fused operators to improve efficiency. Therefore, the Mamba model has corresponding fused operators for Conv operations and recursive formulas, and both GDN and KDA also have corresponding fused operator implementations.
Recursive Formulas: Based on GND, KDA modified the alpha operation part. The FIA (flash-linear-attention) library contains the implementation of the GND Triton version, which currently adds the KDA logic branch:
# Code location: sglang/python/sglang/srt/layers/attention/fla/fused_recurrent.py# Function: def fused_recurrent_gated_delta_rule_fwd_kernel# Added IS_KDA=True, logic adjustment:# ... if not IS_KDA: p_g = g + bos * HV + i_hv else: p_gk = g + (bos * HV + i_hv) * K + o_k# ... if not IS_KDA: b_g = tl.load(p_g).to(tl.float32) b_h *= exp(b_g) else: b_gk = tl.load(p_gk).to(tl.float32) b_h *= exp(b_gk[:, None])# ... if not IS_KDA: p_g += HV else: p_gk += HV * K# Function entry# Code location: sglang/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py# L 386initial_state = ssm_states[cache_indices].contiguous() fused_recurrent_kda( q=q, k=k, v=v, g=g, beta=beta, initial_state=initial_state, use_qk_l2norm_in_kernel=True, cu_seqlens=query_start_loc, )Where the cache acquisition method is:
layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer_id) q_conv_state, k_conv_state, v_conv_state = layer_cache.conv ssm_states = layer_cache.temporalChunkwise Formulas, the input processing is consistent with the previous one, but there are differences in the calling of operators:
# Function entry# Code location: sglang/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py# L 399 initial_state = ssm_states[cache_indices].contiguous() ( core_attn_out, last_recurrent_state, ) = chunk_kda( q=q, k=k, v=v, g=g, beta=beta, initial_state=initial_state, output_final_state=True, use_qk_l2norm_in_kernel=True, cu_seqlens=query_start_loc, )# Operator implementation# Code location: sglang/python/sglang/srt/layers/attention/fla/kda.py# Key function interface: chunk_kda_fwd # L1150def chunk_kda_fwd( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, g: torch.Tensor, beta: torch.Tensor, scale: float, initial_state: torch.Tensor, output_final_state: bool, cu_seqlens: torch.LongTensor | None = None,)2.2 Adaptation of Cache Management
A cache management pool has been defined for Mamba class models in SGLang, and KimiLinear has made modifications based on this:
# Code location: sglang/python/sglang/srt/mem_cache/memory_pool.py# L 124 class MambaPool# Added branch judgment: self.is_kda_cache = isinstance(cache_params, KimiLinearCacheParams)# ... Partial logic: if self.is_kda_cache: conv_state = [ torch.zeros( size=(num_mamba_layers, size + 1) + conv_shape, dtype=conv_dtype, device=device, ) for conv_shape in conv_state_shape ]# ...2.3 Current Related PRs
Kimi Linear
Main PRs:
-
vLLM: https://github.com/vllm-project/vllm/pull/27809/vllm/pull/27809
-
SGLang: https://github.com/sgl-project/sglang/pull/12469
Qwen Next
The implementation of Qwen Next in vLLM/SGLang is similar and also inherits the existing logic of Mamba. Main PRs:
-
SGLang: https://github.com/sgl-project/sglang/pull/10233
-
vLLM: https://github.com/vllm-project/vllm/pull/24526
Key file code locations:
-
Model: vllm/vllm/model_executor/models/qwen3_next.py
-
Operator: vllm/vllm/model_executor/layers/fla/ops/fused_recurrent.py
-
Mamba Base Class: vllm/vllm/model_executor/layers/mamba/abstract.py
It is worth noting that the performance of models in open-source frameworks may be lower than the official data. On one hand, model vendors generally have their own internal versions, which iterate faster, leading to content that is ahead of the open-source version; on the other hand, the hardware used may not be NVIDIA GPUs, while the currently open-source content is mainly targeted at GPUs. To optimize the performance of open-source linear modules, one can consider combining hardware characteristics and focusing on aspects such as operator multi-streaming and graph mode.
Reference Content:
-
In-depth analysis of the KV cache management mechanism in vLLM V1
-
How much advantage does Linear Attention have in KV cache storage? A computational deduction
-
https://arxiv.org/abs/2412.06464
-
https://github.com/MoonshotAI/Kimi-Linear
-
https://qwen.ai/blog?id=4074cca80393150c248e508aa62983f9cb7d27cd&from=research.latest-advancements-list
-
https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct/blob/main/config.json
-
https://huggingface.co/moonshotai/Kimi-Linear-48B-A3B-Instruct/blob/main/config.json
-
https://github.com/vllm-project/vllm
-
https://github.com/sgl-project/sglang
-
https://pytorch.org/blog/hybrid-models-as-first-class-citizens-in-vllm/
-
https://github.com/fla-org/flash-linear-attention
Scan to follow us for more AI Infra foundational knowledge.
-
h-linear-attention