@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 824058076c
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 8, 22:45 UTC

[PERF] Change GDN Attention State Layout from [N, HV, K, V] to [N, HV, V, K] (#33291)

Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>

Vadim Gimpelson committed Feb 4, 2026 at 15:20 UTC 824058076c56164a3772a5f5829bd9662507e5a3
6 files changed +61 -61
vllm/model_executor/layers/fla/ops/chunk.py
+4 -4
@@ -138,11 +138,11 @@ def chunk_gated_delta_rule(
138 Scale factor for the RetNet attention scores.
139 If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
140 initial_state (Optional[torch.Tensor]):
141 - Initial state of shape `[N, H, K, V]` for `N` input sequences.
141 + Initial state of shape `[N, H, V, K]` for `N` input sequences.
142 For equal-length input sequences, `N` equals the batch size `B`.
143 Default: `None`.
144 output_final_state (Optional[bool]):
145 - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`.
145 + Whether to output the final state of shape `[N, H, V, K]`. Default: `False`.
146 cu_seqlens (torch.LongTensor):
147 Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
148 consistent with the FlashAttention API.
@@ -154,7 +154,7 @@ def chunk_gated_delta_rule(
154 o (torch.Tensor):
155 Outputs of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`.
156 final_state (torch.Tensor):
157 - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`.
157 + Final state of shape `[N, H, V, K]` if `output_final_state=True` else `None`.
158
159 Examples::
160 >>> import torch
@@ -168,7 +168,7 @@ def chunk_gated_delta_rule(
168 >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda')
169 >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid()
170 >>> g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda'))
171 - >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda')
171 + >>> h0 = torch.randn(B, H, V, K, dtype=torch.bfloat16, device='cuda')
172 >>> o, ht = chunk_gated_delta_rule(
173 q, k, v, g, beta,
174 initial_state=h0,
vllm/model_executor/layers/fla/ops/chunk_delta_h.py
+35 -35
@@ -81,70 +81,70 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
81 NT = tl.cdiv(T, BT)
82 boh = i_n * NT
83
84 - # [BK, BV]
85 - b_h1 = tl.zeros([64, BV], dtype=tl.float32)
84 + # [BV, BK]
85 + b_h1 = tl.zeros([BV, 64], dtype=tl.float32)
86 if K > 64:
87 - b_h2 = tl.zeros([64, BV], dtype=tl.float32)
87 + b_h2 = tl.zeros([BV, 64], dtype=tl.float32)
88 if K > 128:
89 - b_h3 = tl.zeros([64, BV], dtype=tl.float32)
89 + b_h3 = tl.zeros([BV, 64], dtype=tl.float32)
90 if K > 192:
91 - b_h4 = tl.zeros([64, BV], dtype=tl.float32)
91 + b_h4 = tl.zeros([BV, 64], dtype=tl.float32)
92
93 # calculate offset
94 - h += ((boh * H + i_h) * K * V).to(tl.int64)
94 + h += ((boh * H + i_h) * V * K).to(tl.int64)
95 v += ((bos * H + i_h) * V).to(tl.int64)
96 k += ((bos * Hg + i_h // (H // Hg)) * K).to(tl.int64)
97 w += ((bos * H + i_h) * K).to(tl.int64)
98 if SAVE_NEW_VALUE:
99 v_new += ((bos * H + i_h) * V).to(tl.int64)
100 stride_v = H * V
101 - stride_h = H * K * V
101 + stride_h = H * V * K
102 stride_k = Hg * K
103 stride_w = H * K
104 if USE_INITIAL_STATE:
105 - h0 = h0 + i_nh * K * V
105 + h0 = h0 + i_nh * V * K
106 if STORE_FINAL_STATE:
107 - ht = ht + i_nh * K * V
107 + ht = ht + i_nh * V * K
108
109 # load initial state
110 if USE_INITIAL_STATE:
111 - p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0))
111 + p_h0_1 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0))
112 b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32)
113 if K > 64:
114 p_h0_2 = tl.make_block_ptr(
115 - h0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)
115 + h0, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)
116 )
117 b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32)
118 if K > 128:
119 p_h0_3 = tl.make_block_ptr(
120 - h0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)
120 + h0, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)
121 )
122 b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32)
123 if K > 192:
124 p_h0_4 = tl.make_block_ptr(
125 - h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)
125 + h0, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)
126 )
127 b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32)
128
129 # main recurrence
130 for i_t in range(NT):
131 p_h1 = tl.make_block_ptr(
132 - h + i_t * stride_h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)
132 + h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)
133 )
134 tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1))
135 if K > 64:
136 p_h2 = tl.make_block_ptr(
137 - h + i_t * stride_h, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)
137 + h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)
138 )
139 tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1))
140 if K > 128:
141 p_h3 = tl.make_block_ptr(
142 - h + i_t * stride_h, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)
142 + h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)
143 )
144 tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1))
145 if K > 192:
146 p_h4 = tl.make_block_ptr(
147 - h + i_t * stride_h, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)
147 + h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)
148 )
149 tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1))
150
@@ -152,25 +152,25 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
152 w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0)
153 )
154 b_w = tl.load(p_w, boundary_check=(0, 1))
155 - b_v = tl.dot(b_w, b_h1.to(b_w.dtype))
155 + b_v = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype))
156 if K > 64:
157 p_w = tl.make_block_ptr(
158 w, (T, K), (stride_w, 1), (i_t * BT, 64), (BT, 64), (1, 0)
159 )
160 b_w = tl.load(p_w, boundary_check=(0, 1))
161 - b_v += tl.dot(b_w, b_h2.to(b_w.dtype))
161 + b_v += tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype))
162 if K > 128:
163 p_w = tl.make_block_ptr(
164 w, (T, K), (stride_w, 1), (i_t * BT, 128), (BT, 64), (1, 0)
165 )
166 b_w = tl.load(p_w, boundary_check=(0, 1))
167 - b_v += tl.dot(b_w, b_h3.to(b_w.dtype))
167 + b_v += tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype))
168 if K > 192:
169 p_w = tl.make_block_ptr(
170 w, (T, K), (stride_w, 1), (i_t * BT, 192), (BT, 64), (1, 0)
171 )
172 b_w = tl.load(p_w, boundary_check=(0, 1))
173 - b_v += tl.dot(b_w, b_h4.to(b_w.dtype))
173 + b_v += tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype))
174 p_v = tl.make_block_ptr(
175 v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)
176 )
@@ -207,7 +207,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
207 mask=(o_k1 < K),
208 other=0.0,
209 )
210 - b_h1 *= exp(b_gk_last1)[:, None]
210 + b_h1 *= exp(b_gk_last1)[None, :]
211 if K > 64:
212 o_k2 = 64 + o_k1
213 b_gk_last2 = tl.load(
@@ -215,7 +215,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
215 mask=(o_k2 < K),
216 other=0.0,
217 )
218 - b_h2 *= exp(b_gk_last2)[:, None]
218 + b_h2 *= exp(b_gk_last2)[None, :]
219 if K > 128:
220 o_k3 = 128 + o_k1
221 b_gk_last3 = tl.load(
@@ -223,7 +223,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
223 mask=(o_k3 < K),
224 other=0.0,
225 )
226 - b_h3 *= exp(b_gk_last3)[:, None]
226 + b_h3 *= exp(b_gk_last3)[None, :]
227 if K > 192:
228 o_k4 = 192 + o_k1
229 b_gk_last4 = tl.load(
@@ -231,49 +231,49 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
231 mask=(o_k4 < K),
232 other=0.0,
233 )
234 - b_h4 *= exp(b_gk_last4)[:, None]
234 + b_h4 *= exp(b_gk_last4)[None, :]
235 b_v = b_v.to(k.dtype.element_ty)
236
237 p_k = tl.make_block_ptr(
238 k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)
239 )
240 b_k = tl.load(p_k, boundary_check=(0, 1))
241 - b_h1 += tl.dot(b_k, b_v)
241 + b_h1 += tl.trans(tl.dot(b_k, b_v))
242 if K > 64:
243 p_k = tl.make_block_ptr(
244 k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)
245 )
246 b_k = tl.load(p_k, boundary_check=(0, 1))
247 - b_h2 += tl.dot(b_k, b_v)
247 + b_h2 += tl.trans(tl.dot(b_k, b_v))
248 if K > 128:
249 p_k = tl.make_block_ptr(
250 k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)
251 )
252 b_k = tl.load(p_k, boundary_check=(0, 1))
253 - b_h3 += tl.dot(b_k, b_v)
253 + b_h3 += tl.trans(tl.dot(b_k, b_v))
254 if K > 192:
255 p_k = tl.make_block_ptr(
256 k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)
257 )
258 b_k = tl.load(p_k, boundary_check=(0, 1))
259 - b_h4 += tl.dot(b_k, b_v)
259 + b_h4 += tl.trans(tl.dot(b_k, b_v))
260 # epilogue
261 if STORE_FINAL_STATE:
262 - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0))
262 + p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0))
263 tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
264 if K > 64:
265 p_ht = tl.make_block_ptr(
266 - ht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)
266 + ht, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)
267 )
268 tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
269 if K > 128:
270 p_ht = tl.make_block_ptr(
271 - ht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)
271 + ht, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)
272 )
273 tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
274 if K > 192:
275 p_ht = tl.make_block_ptr(
276 - ht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)
276 + ht, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)
277 )
278 tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
279
@@ -312,9 +312,9 @@ def chunk_gated_delta_rule_fwd_h(
312 )
313 assert K <= 256, "current kernel does not support head dimension larger than 256."
314
315 - h = k.new_empty(B, NT, H, K, V)
315 + h = k.new_empty(B, NT, H, V, K)
316 final_state = (
317 - k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None
317 + k.new_empty(N, H, V, K, dtype=torch.float32) if output_final_state else None
318 )
319
320 v_new = torch.empty_like(u) if save_new_value else None
vllm/model_executor/layers/fla/ops/chunk_o.py
+4 -4
@@ -85,7 +85,7 @@ def chunk_fwd_kernel_o(
85 k += (bos * Hg + i_h // (H // Hg)) * K
86 v += (bos * H + i_h) * V
87 o += (bos * H + i_h) * V
88 - h += (i_tg * H + i_h).to(tl.int64) * K * V
88 + h += (i_tg * H + i_h).to(tl.int64) * V * K
89
90 b_o = tl.zeros([BT, BV], dtype=tl.float32)
91 b_A = tl.zeros([BT, BT], dtype=tl.float32)
@@ -98,17 +98,17 @@ def chunk_fwd_kernel_o(
98 k, (K, T), (1, Hg * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)
99 )
100 p_h = tl.make_block_ptr(
101 - h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)
101 + h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)
102 )
103 # [BT, BK]
104 b_q = tl.load(p_q, boundary_check=(0, 1))
105 # [BK, BT]
106 b_k = tl.load(p_k, boundary_check=(0, 1))
107 - # [BK, BV]
107 + # [BV, BK]
108 b_h = tl.load(p_h, boundary_check=(0, 1))
109
110 # [BT, BK] @ [BK, BV] -> [BT, BV]
111 - b_o += tl.dot(b_q, b_h)
111 + b_o += tl.dot(b_q, tl.trans(b_h))
112 # [BT, BK] @ [BK, BT] -> [BT, BT]
113 b_A += tl.dot(b_q, b_k)
114
vllm/model_executor/layers/fla/ops/fused_recurrent.py
+16 -16
@@ -97,9 +97,9 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
97
98 mask_k = o_k < K
99 mask_v = o_v < V
100 - mask_h = mask_k[:, None] & mask_v[None, :]
100 + mask_h = mask_v[:, None] & mask_k[None, :]
101
102 - b_h = tl.zeros([BK, BV], dtype=tl.float32)
102 + b_h = tl.zeros([BV, BK], dtype=tl.float32)
103 if USE_INITIAL_STATE:
104 if IS_CONTINUOUS_BATCHING:
105 if IS_SPEC_DECODING:
@@ -115,8 +115,8 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
115 return
116 p_h0 = h0 + state_idx * stride_init_state_token
117 else:
118 - p_h0 = h0 + bos * HV * K * V
119 - p_h0 = p_h0 + i_hv * K * V + o_k[:, None] * V + o_v[None, :]
118 + p_h0 = h0 + bos * HV * V * K
119 + p_h0 = p_h0 + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
120 b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
121
122 for i_t in range(0, T):
@@ -128,24 +128,24 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
128 b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)
129 b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6)
130 b_q = b_q * scale
131 - # [BK, BV]
131 + # [BV, BK]
132 if not IS_KDA:
133 b_g = tl.load(p_g).to(tl.float32)
134 b_h *= exp(b_g)
135 else:
136 b_gk = tl.load(p_gk).to(tl.float32)
137 - b_h *= exp(b_gk[:, None])
137 + b_h *= exp(b_gk[None, :])
138 # [BV]
139 - b_v -= tl.sum(b_h * b_k[:, None], 0)
139 + b_v -= tl.sum(b_h * b_k[None, :], 1)
140 if IS_BETA_HEADWISE:
141 b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32)
142 else:
143 b_beta = tl.load(p_beta).to(tl.float32)
144 b_v *= b_beta
145 - # [BK, BV]
146 - b_h += b_k[:, None] * b_v[None, :]
145 + # [BV, BK]
146 + b_h += b_v[:, None] * b_k[None, :]
147 # [BV]
148 - b_o = tl.sum(b_h * b_q[:, None], 0)
148 + b_o = tl.sum(b_h * b_q[None, :], 1)
149 tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
150
151 # keep the states for multi-query tokens
@@ -157,11 +157,11 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
157 # Only store if state index is valid (not PAD_SLOT_ID)
158 if final_state_idx >= 0:
159 p_ht = ht + final_state_idx * stride_final_state_token
160 - p_ht = p_ht + i_hv * K * V + o_k[:, None] * V + o_v[None, :]
160 + p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
161 tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
162 else:
163 p_ht = ht + (bos + i_t) * stride_final_state_token
164 - p_ht = p_ht + i_hv * K * V + o_k[:, None] * V + o_v[None, :]
164 + p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
165 tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
166
167 p_q += H * K
@@ -202,7 +202,7 @@ def fused_recurrent_gated_delta_rule_fwd(
202 if inplace_final_state:
203 final_state = initial_state
204 else:
205 - final_state = q.new_empty(T, HV, K, V, dtype=initial_state.dtype)
205 + final_state = q.new_empty(T, HV, V, K, dtype=initial_state.dtype)
206
207 stride_init_state_token = initial_state.stride(0)
208 stride_final_state_token = final_state.stride(0)
@@ -318,7 +318,7 @@ def fused_recurrent_gated_delta_rule(
318 Scale factor for the RetNet attention scores.
319 If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
320 initial_state (Optional[torch.Tensor]):
321 - Initial state of shape `[N, HV, K, V]` for `N` input sequences.
321 + Initial state of shape `[N, HV, V, K]` for `N` input sequences.
322 For equal-length input sequences, `N` equals the batch size `B`.
323 Default: `None`.
324 inplace_final_state: bool:
@@ -336,7 +336,7 @@ def fused_recurrent_gated_delta_rule(
336 o (torch.Tensor):
337 Outputs of shape `[B, T, HV, V]`.
338 final_state (torch.Tensor):
339 - Final state of shape `[N, HV, K, V]`.
339 + Final state of shape `[N, HV, V, K]`.
340
341 Examples::
342 >>> import torch
@@ -350,7 +350,7 @@ def fused_recurrent_gated_delta_rule(
350 >>> v = torch.randn(B, T, HV, V, device='cuda')
351 >>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda'))
352 >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid()
353 - >>> h0 = torch.randn(B, HV, K, V, device='cuda')
353 + >>> h0 = torch.randn(B, HV, V, K, device='cuda')
354 >>> o, ht = fused_gated_recurrent_delta_rule(
355 q, k, v, g, beta,
356 initial_state=h0,
vllm/model_executor/layers/fla/ops/kda.py
+1 -1
@@ -55,7 +55,7 @@ def fused_recurrent_kda_fwd(
55 if inplace_final_state:
56 final_state = initial_state
57 else:
58 - final_state = q.new_empty(T, HV, K, V, dtype=initial_state.dtype)
58 + final_state = q.new_empty(T, HV, V, K, dtype=initial_state.dtype)
59
60 stride_init_state_token = initial_state.stride(0)
61 stride_final_state_token = final_state.stride(0)
vllm/model_executor/layers/mamba/mamba_utils.py
+1 -1
@@ -191,8 +191,8 @@ class MambaStateShapeCalculator:
191
192 temporal_state_shape = (
193 divide(num_v_heads, tp_world_size),
194 - head_k_dim,
194 head_v_dim,
195 + head_k_dim,
196 )
197 return conv_state_shape, temporal_state_shape
198