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 14, 00:30 UTC
[Perf] Decouple torch op from GDA to leverage torch.compile (#27871)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Jiangyun Zhu committed
Oct 31, 2025 at 21:35 UTC
3857eb87257cf23d2401a651884dcbbf42c16b7b
1 file changed
+68
-48
vllm/model_executor/layers/kda.py
+68
-48
@@ -40,18 +40,36 @@ logger = init_logger(__name__)
40
41
42
def kda_attention(
43
- hidden_states: torch.Tensor,
44
- output: torch.Tensor,
43
+ q_proj_states: torch.Tensor,
44
+ k_proj_states: torch.Tensor,
45
+ v_proj_states: torch.Tensor,
46
+ g1: torch.Tensor,
47
+ g2: torch.Tensor,
48
+ beta: torch.Tensor,
49
+ core_attn_out: torch.Tensor,
50
layer_name: str,
51
) -> None:
52
forward_context: ForwardContext = get_forward_context()
53
self = forward_context.no_compile_layers[layer_name]
49
- self._forward(hidden_states=hidden_states, output=output)
54
+ self._forward(
55
+ q_proj_states=q_proj_states,
56
+ k_proj_states=k_proj_states,
57
+ v_proj_states=v_proj_states,
58
+ g1=g1,
59
+ g2=g2,
60
+ beta=beta,
61
+ core_attn_out=core_attn_out,
62
+ )
63
64
65
def kda_attention_fake(
53
- hidden_states: torch.Tensor,
54
- output: torch.Tensor,
66
+ q_proj_states: torch.Tensor,
67
+ k_proj_states: torch.Tensor,
68
+ v_proj_states: torch.Tensor,
69
+ g1: torch.Tensor,
70
+ g2: torch.Tensor,
71
+ beta: torch.Tensor,
72
+ core_attn_out: torch.Tensor,
73
layer_name: str,
74
) -> None:
75
return
@@ -60,7 +78,7 @@ def kda_attention_fake(
78
direct_register_custom_op(
79
op_name="kda_attention",
80
op_func=kda_attention,
63
- mutates_args=["output"],
81
+ mutates_args=["core_attn_out"],
82
fake_impl=kda_attention_fake,
83
)
84
@@ -241,37 +259,56 @@ class KimiDeltaAttention(nn.Module, MambaBase):
259
hidden_states: torch.Tensor,
260
positions: torch.Tensor,
261
output: torch.Tensor,
244
- ) -> None:
245
- return torch.ops.vllm.kda_attention(
246
- hidden_states,
247
- output,
262
+ ) -> torch.Tensor:
263
+ num_tokens = hidden_states.size(0)
264
+ q = self.q_proj(hidden_states)[0]
265
+ k = self.k_proj(hidden_states)[0]
266
+ v = self.v_proj(hidden_states)[0]
267
+
268
+ beta = self.b_proj(hidden_states)[0].float().sigmoid()
269
+ g1 = self.f_b_proj(self.f_a_proj(hidden_states)[0])[0]
270
+ g1 = fused_kda_gate(g1, self.A_log, self.head_dim, g_bias=self.dt_bias)
271
+ beta = beta.unsqueeze(0)
272
+ g1 = g1.unsqueeze(0)
273
+
274
+ g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0]
275
+ g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim)
276
+
277
+ core_attn_out = torch.zeros(
278
+ (1, num_tokens, self.local_num_heads, self.head_dim),
279
+ dtype=hidden_states.dtype,
280
+ device=hidden_states.device,
281
+ )
282
+ torch.ops.vllm.kda_attention(
283
+ q,
284
+ k,
285
+ v,
286
+ g1,
287
+ g2,
288
+ beta,
289
+ core_attn_out,
290
self.prefix,
291
)
292
+ core_attn_out = self.o_norm(core_attn_out, g2)
293
+ core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
294
+
295
+ return self.o_proj(core_attn_out)[0]
296
297
def _forward(
298
self,
253
- hidden_states: torch.Tensor,
254
- output: torch.Tensor,
299
+ q_proj_states: torch.Tensor,
300
+ k_proj_states: torch.Tensor,
301
+ v_proj_states: torch.Tensor,
302
+ g1: torch.Tensor,
303
+ g2: torch.Tensor,
304
+ beta: torch.Tensor,
305
+ core_attn_out: torch.Tensor,
306
) -> None:
307
forward_context = get_forward_context()
308
attn_metadata: AttentionMetadata = forward_context.attn_metadata
309
310
if attn_metadata is None:
260
- # V1 profile run
261
- # Mimic the memory allocation in the real run
262
- q = torch.empty_like(hidden_states)
263
- k = torch.empty_like(hidden_states)
264
- v = torch.empty_like(hidden_states)
265
- g = hidden_states.new_empty(
266
- hidden_states.size(0),
267
- self.local_num_heads,
268
- self.head_dim,
269
- dtype=torch.float32,
270
- )
271
- beta = torch.empty(
272
- hidden_states.size(0), self.local_num_heads, dtype=torch.float32
273
- )
274
- core_attn_out = torch.empty_like(hidden_states)
311
+ # # V1 profile run
312
return
313
314
assert isinstance(attn_metadata, dict)
@@ -288,10 +325,6 @@ class KimiDeltaAttention(nn.Module, MambaBase):
325
conv_state_k = conv_state_k.transpose(-1, -2)
326
conv_state_v = conv_state_v.transpose(-1, -2)
327
291
- q_proj_states = self.q_proj(hidden_states)[0]
292
- k_proj_states = self.k_proj(hidden_states)[0]
293
- v_proj_states = self.v_proj(hidden_states)[0]
294
-
328
q_conv_weights = self.q_conv1d.weight.view(
329
self.q_conv1d.weight.size(0), self.q_conv1d.weight.size(2)
330
)
@@ -374,14 +407,6 @@ class KimiDeltaAttention(nn.Module, MambaBase):
407
lambda x: rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim), (q, k, v)
408
)
409
377
- beta = self.b_proj(hidden_states)[0].float().sigmoid()
378
-
379
- g = self.f_b_proj(self.f_a_proj(hidden_states)[0])[0]
380
- g = fused_kda_gate(g, self.A_log, self.head_dim, g_bias=self.dt_bias)
381
-
382
- beta = beta.unsqueeze(0)
383
- g = g.unsqueeze(0)
384
-
410
if attn_metadata.num_prefills > 0:
411
zero_idx = non_spec_state_indices_tensor[~has_initial_state]
412
recurrent_state[zero_idx] = 0
@@ -393,7 +418,7 @@ class KimiDeltaAttention(nn.Module, MambaBase):
418
q=q,
419
k=k,
420
v=v,
396
- g=g,
421
+ g=g1,
422
beta=beta,
423
initial_state=initial_state,
424
output_final_state=True,
@@ -410,17 +435,12 @@ class KimiDeltaAttention(nn.Module, MambaBase):
435
q=q,
436
k=k,
437
v=v,
413
- g=g,
438
+ g=g1,
439
beta=beta,
440
initial_state=recurrent_state,
441
use_qk_l2norm_in_kernel=True,
442
cu_seqlens=non_spec_query_start_loc,
443
ssm_state_indices=non_spec_state_indices_tensor,
444
)
420
-
421
- g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0]
422
- g = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim)
423
- core_attn_out = self.o_norm(core_attn_out_non_spec, g)
424
- core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
425
-
426
- output[:] = self.o_proj(core_attn_out)[0]
445
+ assert core_attn_out_non_spec.shape == core_attn_out.shape
446
+ core_attn_out[:] = core_attn_out_non_spec