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, 05:30 UTC
[Kernel] Support DCP for Triton backend (#25132)
Signed-off-by: Wei Wei <wwei6@meta.com>
Wei Wei committed
Sep 24, 2025 at 18:09 UTC
05c19485a5294fef5bbab625822ff96613cf0c39
4 files changed
+30
-8
tests/kernels/attention/test_triton_decode_attention.py
+5
@@ -46,6 +46,8 @@ def test_decode_attention(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE):
46
# o will have the same shape as q
47
o = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda")
48
49
+ lse = torch.zeros(B, H_Q, dtype=dtype, device="cuda")
50
+
51
b_seq_len = torch.full((B, ), seq_len, device="cuda")
52
53
attn_logits = torch.empty(
@@ -60,6 +62,7 @@ def test_decode_attention(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE):
62
k_buffer,
63
v_buffer,
64
o,
65
+ lse,
66
req_to_token,
67
b_seq_len,
68
attn_logits,
@@ -72,12 +75,14 @@ def test_decode_attention(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE):
75
v_buffer = v_buffer.view(CACHE_SIZE // PAGE_SIZE, PAGE_SIZE, H_KV, D_V)
76
77
o1 = torch.zeros_like(o)
78
+ lse1 = torch.zeros_like(lse)
79
80
decode_attention_fwd(
81
q,
82
k_buffer,
83
v_buffer,
84
o1,
85
+ lse1,
86
req_to_page,
87
b_seq_len,
88
attn_logits,
vllm/attention/ops/triton_decode_attention.py
+17
-2
@@ -474,12 +474,14 @@ def _decode_grouped_att_m_fwd(
474
def _fwd_kernel_stage2(
475
Mid_O,
476
o,
477
+ lse,
478
B_Seqlen,
479
stride_mid_ob,
480
stride_mid_oh,
481
stride_mid_os,
482
stride_obs,
483
stride_oh,
484
+ stride_lse_bs,
485
NUM_KV_SPLITS: tl.constexpr,
486
BLOCK_DV: tl.constexpr,
487
Lv: tl.constexpr,
@@ -525,12 +527,18 @@ def _fwd_kernel_stage2(
527
acc / e_sum,
528
mask=mask_d,
529
)
530
+ lse_val = e_max + tl.log(e_sum)
531
+ tl.store(
532
+ lse + cur_batch * stride_lse_bs + cur_head,
533
+ lse_val,
534
+ )
535
536
537
def _decode_softmax_reducev_fwd(
538
logits,
539
q,
540
o,
541
+ lse,
542
v_buffer,
543
b_seq_len,
544
num_kv_splits,
@@ -555,12 +563,14 @@ def _decode_softmax_reducev_fwd(
563
_fwd_kernel_stage2[grid](
564
logits,
565
o,
566
+ lse,
567
b_seq_len,
568
logits.stride(0),
569
logits.stride(1),
570
logits.stride(2),
571
o.stride(0),
572
o.stride(1),
573
+ lse.stride(0),
574
NUM_KV_SPLITS=NUM_KV_SPLITS,
575
BLOCK_DV=BLOCK_DV,
576
Lv=Lv,
@@ -575,6 +585,7 @@ def decode_attention_fwd_normal(
585
k_buffer,
586
v_buffer,
587
o,
588
+ lse,
589
req_to_token,
590
b_seq_len,
591
attn_logits,
@@ -595,7 +606,7 @@ def decode_attention_fwd_normal(
606
page_size,
607
logit_cap,
608
)
598
- _decode_softmax_reducev_fwd(attn_logits, q, o, v_buffer, b_seq_len,
609
+ _decode_softmax_reducev_fwd(attn_logits, q, o, lse, v_buffer, b_seq_len,
610
num_kv_splits)
611
612
@@ -604,6 +615,7 @@ def decode_attention_fwd_grouped(
615
k_buffer,
616
v_buffer,
617
o,
618
+ lse,
619
req_to_token,
620
b_seq_len,
621
attn_logits,
@@ -624,7 +636,7 @@ def decode_attention_fwd_grouped(
636
page_size,
637
logit_cap,
638
)
627
- _decode_softmax_reducev_fwd(attn_logits, q, o, v_buffer, b_seq_len,
639
+ _decode_softmax_reducev_fwd(attn_logits, q, o, lse, v_buffer, b_seq_len,
640
num_kv_splits)
641
642
@@ -633,6 +645,7 @@ def decode_attention_fwd(
645
k_buffer,
646
v_buffer,
647
o,
648
+ lse,
649
req_to_token,
650
b_seq_len,
651
attn_logits,
@@ -651,6 +664,7 @@ def decode_attention_fwd(
664
k_buffer,
665
v_buffer,
666
o,
667
+ lse,
668
req_to_token,
669
b_seq_len,
670
attn_logits,
@@ -666,6 +680,7 @@ def decode_attention_fwd(
680
k_buffer,
681
v_buffer,
682
o,
683
+ lse,
684
req_to_token,
685
b_seq_len,
686
attn_logits,
vllm/model_executor/models/deepseek_v2.py
+1
-1
@@ -685,7 +685,7 @@ class DeepseekV2DecoderLayer(nn.Module):
685
) -> torch.Tensor:
686
# Self Attention
687
if residual is None:
688
- residual = hidden_states
688
+ residual = hidden_states.clone()
689
hidden_states = self.input_layernorm(hidden_states)
690
else:
691
hidden_states, residual = self.input_layernorm(
vllm/v1/attention/backends/mla/triton_mla.py
+7
-5
@@ -32,6 +32,7 @@ class TritonMLABackend(MLACommonBackend):
32
33
34
class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
35
+ can_return_lse_for_decode: bool = True
36
37
def __init__(
38
self,
@@ -139,19 +140,20 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
140
141
assert isinstance(q, torch.Tensor)
142
B = q.shape[0]
143
+ q_num_heads = q.shape[1]
144
o = torch.zeros(B,
143
- self.num_heads,
145
+ q_num_heads,
146
self.kv_lora_rank,
147
dtype=q.dtype,
148
device=q.device)
147
-
149
+ lse = torch.zeros(B, q_num_heads, dtype=q.dtype, device=q.device)
150
num_kv_splits = 4 # TODO: heuristic
151
152
# TODO(lucas) Allocate ahead of time
153
attn_logits = torch.empty(
154
(
155
B,
154
- self.num_heads,
156
+ q_num_heads,
157
num_kv_splits,
158
# NOTE(lucas) idk why the +1 is here but sglang has it so we
159
# just mirror that
@@ -167,9 +169,9 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
169
PAGE_SIZE = kv_c_and_k_pe_cache.size(1)
170
171
# Run MQA
170
- decode_attention_fwd(q, kv_c_and_k_pe_cache, kv_c_cache, o,
172
+ decode_attention_fwd(q, kv_c_and_k_pe_cache, kv_c_cache, o, lse,
173
attn_metadata.decode.block_table,
174
attn_metadata.decode.seq_lens, attn_logits,
175
num_kv_splits, self.scale, PAGE_SIZE)
176
175
- return o, None
177
+ return o, lse