[Perf] Add fused DFlash2 grouped convolution (#55960)
Signed-off-by: mgoin <mgoin64@gmail.com>
Misha Goin committed
Sep 17, 2026 at 17:38 UTC
db7a24c230a4db6ae3a568ed01b01f3f7172069d
4 files changed
+199
.buildkite/test_areas/spec_decode.yaml
+4
@@ -283,6 +283,8 @@ steps:
283
- vllm/v1/spec_decode/
284
- vllm/v1/worker/gpu/spec_decode/
285
- vllm/model_executor/models/qwen3_dflash.py
286
+ - vllm/model_executor/models/qwen3_dflash2.py
287
+ - vllm/model_executor/models/qwen3_5.py
288
- vllm/model_executor/models/laguna_dflash.py
289
- tests/v1/e2e/spec_decode/
290
commands:
@@ -293,6 +295,8 @@ steps:
295
dind: false
296
device: mi300_1
297
timeout_in_minutes: 120
298
+ commands:
299
+ - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ -k "not dflash2"
300
depends_on:
301
- image-build-amd
302
source_file_dependencies:
tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py
+21
@@ -26,6 +26,8 @@ class DFlashCorrectnessConfig:
26
use_chat_completions: bool = False
27
enforce_eager: bool = False
28
disable_flashinfer_sampler: bool = False
29
+ language_model_only: bool = False
30
+ chat_template_kwargs: dict[str, object] | None = None
31
32
33
QWEN3_DFLASH = DFlashCorrectnessConfig(
@@ -35,6 +37,18 @@ QWEN3_DFLASH = DFlashCorrectnessConfig(
37
expected_acceptance_len=3.5,
38
)
39
40
+QWEN3_8_DFLASH2_NVFP4 = DFlashCorrectnessConfig(
41
+ model="nvidia/Qwen3.8-27B-NVFP4",
42
+ draft_model="z-lab/Qwen3.8-27B-DFlash2",
43
+ expected_accuracy=0.90,
44
+ expected_acceptance_len=6.14 * 0.95,
45
+ num_speculative_tokens=7,
46
+ max_num_seqs=32,
47
+ use_chat_completions=True,
48
+ language_model_only=True,
49
+ chat_template_kwargs={"enable_thinking": False},
50
+)
51
+
52
LAGUNA_DFLASH_NVFP4 = DFlashCorrectnessConfig(
53
model="poolside/Laguna-XS-2.1-NVFP4",
54
draft_model="poolside/Laguna-XS-2.1-DFlash-NVFP4",
@@ -104,6 +118,11 @@ def test_dflash_reference_acceptance_lengths(
118
True,
119
id="laguna-nvfp4-mrv2",
120
),
121
+ pytest.param(
122
+ QWEN3_8_DFLASH2_NVFP4,
123
+ True,
124
+ id="qwen3.8-dflash2-nvfp4-mrv2",
125
+ ),
126
],
127
)
128
def test_dflash_correctness(
@@ -134,12 +153,14 @@ def test_dflash_correctness(
153
disable_log_stats=False,
154
enable_chunked_prefill=None,
155
compilation_config=CompilationConfig(),
156
+ language_model_only=config.language_model_only,
157
) as spec_runner:
158
spec_llm = spec_runner.llm
159
results = evaluate_gsm8k_offline(
160
spec_llm,
161
num_questions=config.num_questions,
162
use_chat_completions=config.use_chat_completions,
163
+ chat_template_kwargs=config.chat_template_kwargs,
164
)
165
accuracy = results["accuracy"]
166
acceptance_len = compute_acceptance_len(spec_llm.get_metrics())
tests/v1/spec_decode/test_dflash2.py
+50
@@ -7,6 +7,7 @@ import pytest
7
import torch
8
9
from vllm.model_executor.models.qwen3_dflash2 import _grouped_conv, _score_edges
10
+from vllm.platforms import current_platform
11
from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
12
from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import DFlash2Speculator
13
@@ -35,6 +36,55 @@ def test_grouped_conv_matches_reference(block_size: int):
36
torch.testing.assert_close(actual, expected.flatten(0, 1).flatten(-2))
37
38
39
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
40
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
41
+@pytest.mark.parametrize(
42
+ "batch,num_groups,group_size,block_size,taps",
43
+ [
44
+ (7, 13, 11, 5, 1),
45
+ (7, 13, 11, 5, 3),
46
+ (7, 13, 11, 8, 2),
47
+ (16, 64, 16, 8, 2),
48
+ (16, 160, 16, 8, 2),
49
+ ],
50
+)
51
+def test_grouped_conv_triton_matches_reference(
52
+ dtype: torch.dtype,
53
+ batch: int,
54
+ num_groups: int,
55
+ group_size: int,
56
+ block_size: int,
57
+ taps: int,
58
+):
59
+ torch.manual_seed(0)
60
+ rows = batch * block_size
61
+ hidden = torch.randn(rows, num_groups * group_size, device="cuda", dtype=dtype)
62
+ base = torch.randn(taps, num_groups * group_size, device="cuda", dtype=dtype)
63
+ projected = torch.randn(rows, 2, taps, num_groups, device="cuda", dtype=dtype)
64
+ delta = projected[:, 1]
65
+
66
+ actual = _grouped_conv(
67
+ hidden, delta, base, block_size, num_groups, group_size, taps
68
+ )
69
+
70
+ hidden_blocks = hidden.float().view(batch, block_size, num_groups, group_size)
71
+ expected = torch.zeros_like(hidden_blocks)
72
+ base_blocks = base.float().view(taps, num_groups, group_size)
73
+ delta_blocks = delta.float().view(batch, block_size, taps, num_groups)
74
+ for position in range(block_size):
75
+ for tap in range(min(taps, position + 1)):
76
+ expected[:, position] += (
77
+ base_blocks[tap] + delta_blocks[:, position, tap, :, None]
78
+ ) * hidden_blocks[:, position - tap]
79
+
80
+ torch.testing.assert_close(
81
+ actual,
82
+ expected.flatten(0, 1).flatten(-2).to(dtype),
83
+ rtol=1e-2 if dtype is torch.bfloat16 else 1e-5,
84
+ atol=1e-2 if dtype is torch.bfloat16 else 1e-5,
85
+ )
86
+
87
+
88
def test_selector_edges_match_sequential_reference():
89
torch.manual_seed(1)
90
batch, steps, top_k, rank = 2, 4, 3, 5
vllm/model_executor/models/qwen3_dflash2.py
+124
@@ -11,6 +11,8 @@ from vllm.config import CacheConfig, VllmConfig
11
from vllm.model_executor.layers.linear import ReplicatedLinear
12
from vllm.model_executor.layers.logits_processor import LogitsProcessor
13
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
14
+from vllm.triton_utils import tl, triton
15
+from vllm.utils.torch_utils import direct_register_custom_op
16
17
from .qwen3_dflash import (
18
DFlashQwen3DecoderLayer,
@@ -20,6 +22,125 @@ from .qwen3_dflash import (
22
from .utils import maybe_prefix
23
24
25
+@triton.jit
26
+def _dflash2_grouped_conv_kernel(
27
+ x_ptr,
28
+ delta_ptr,
29
+ base_ptr,
30
+ output_ptr,
31
+ x_stride_row,
32
+ delta_stride_row,
33
+ delta_stride_tap,
34
+ base_stride_tap,
35
+ output_stride_row,
36
+ NUM_CHANNELS: tl.constexpr,
37
+ BLOCK_SIZE: tl.constexpr,
38
+ GROUP_SIZE: tl.constexpr,
39
+ TAPS: tl.constexpr,
40
+ ELEMENT_BLOCK: tl.constexpr,
41
+) -> None:
42
+ row = tl.program_id(0) // triton.cdiv(NUM_CHANNELS, ELEMENT_BLOCK)
43
+ col_block = tl.program_id(0) % triton.cdiv(NUM_CHANNELS, ELEMENT_BLOCK)
44
+ channels = col_block * ELEMENT_BLOCK + tl.arange(0, ELEMENT_BLOCK)
45
+ mask = channels < NUM_CHANNELS
46
+ groups = channels // GROUP_SIZE
47
+ position = row % BLOCK_SIZE
48
+
49
+ delta_row = delta_ptr + row * delta_stride_row
50
+ x_row = x_ptr + row * x_stride_row
51
+ accumulator = (
52
+ tl.load(base_ptr + channels, mask=mask, other=0.0).to(tl.float32)
53
+ + tl.load(delta_row + groups, mask=mask, other=0.0).to(tl.float32)
54
+ ) * tl.load(x_row + channels, mask=mask, other=0.0).to(tl.float32)
55
+
56
+ for tap in tl.static_range(1, TAPS):
57
+ tap_mask = mask & (position >= tap)
58
+ coefficient = tl.load(
59
+ base_ptr + tap * base_stride_tap + channels,
60
+ mask=tap_mask,
61
+ other=0.0,
62
+ ).to(tl.float32) + tl.load(
63
+ delta_row + tap * delta_stride_tap + groups,
64
+ mask=tap_mask,
65
+ other=0.0,
66
+ ).to(tl.float32)
67
+ x = tl.load(
68
+ x_ptr + (row - tap) * x_stride_row + channels,
69
+ mask=tap_mask,
70
+ other=0.0,
71
+ ).to(tl.float32)
72
+ accumulator += coefficient * x
73
+
74
+ tl.store(
75
+ output_ptr + row * output_stride_row + channels,
76
+ accumulator,
77
+ mask=mask,
78
+ )
79
+
80
+
81
+def dflash2_grouped_conv_impl(
82
+ x: torch.Tensor,
83
+ delta: torch.Tensor,
84
+ base: torch.Tensor,
85
+ block_size: int,
86
+ group_size: int,
87
+) -> torch.Tensor:
88
+ """DFlash2 grouped convolution adapted from LightSeek TokenSpeed PR #1399."""
89
+ num_rows, num_channels = x.shape
90
+ output = torch.empty_like(x)
91
+ if num_rows == 0:
92
+ return output
93
+
94
+ element_block = 1024 if num_rows >= 128 and num_channels % 1024 == 0 else 512
95
+ grid = (num_rows * triton.cdiv(num_channels, element_block),)
96
+ _dflash2_grouped_conv_kernel[grid](
97
+ x,
98
+ delta,
99
+ base,
100
+ output,
101
+ x.stride(0),
102
+ delta.stride(0),
103
+ delta.stride(1),
104
+ base.stride(0),
105
+ output.stride(0),
106
+ NUM_CHANNELS=num_channels,
107
+ BLOCK_SIZE=block_size,
108
+ GROUP_SIZE=group_size,
109
+ TAPS=base.shape[0],
110
+ ELEMENT_BLOCK=element_block,
111
+ num_warps=4,
112
+ )
113
+ return output
114
+
115
+
116
+def dflash2_grouped_conv_fake(
117
+ x: torch.Tensor,
118
+ delta: torch.Tensor,
119
+ base: torch.Tensor,
120
+ block_size: int,
121
+ group_size: int,
122
+) -> torch.Tensor:
123
+ return torch.empty_like(x)
124
+
125
+
126
+direct_register_custom_op(
127
+ op_name="dflash2_grouped_conv",
128
+ op_func=dflash2_grouped_conv_impl,
129
+ fake_impl=dflash2_grouped_conv_fake,
130
+ tags=(torch.Tag.needs_fixed_stride_order,),
131
+)
132
+
133
+
134
+def dflash2_grouped_conv(
135
+ x: torch.Tensor,
136
+ delta: torch.Tensor,
137
+ base: torch.Tensor,
138
+ block_size: int,
139
+ group_size: int,
140
+) -> torch.Tensor:
141
+ return torch.ops.vllm.dflash2_grouped_conv(x, delta, base, block_size, group_size)
142
+
143
+
144
def _grouped_conv(
145
hidden_states: torch.Tensor,
146
delta: torch.Tensor,
@@ -29,6 +150,9 @@ def _grouped_conv(
150
group_size: int,
151
taps: int,
152
) -> torch.Tensor:
153
+ if hidden_states.is_cuda:
154
+ return dflash2_grouped_conv(hidden_states, delta, base, block_size, group_size)
155
+
156
blocks = hidden_states.unflatten(-1, (num_groups, group_size))
157
coefficients = base.view(1, taps, num_groups, group_size) + delta.unsqueeze(-1)
158
output = coefficients[:, 0] * blocks