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, 01:15 UTC
[Refactor] Deprecate `head_first` for `chunk_gated_delta_rule` (#34263)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Feb 19, 2026 at 13:23 UTC
c683d11c94655655cd7bf95a27aef7e245325102
3 files changed
+8
-36
vllm/model_executor/layers/fla/ops/chunk.py
+8
-29
@@ -10,7 +10,6 @@
10
import warnings
11
12
import torch
13
-from einops import rearrange
13
14
from .chunk_delta_h import chunk_gated_delta_rule_fwd_h
15
from .chunk_o import chunk_fwd_o
@@ -119,21 +118,20 @@ def chunk_gated_delta_rule(
118
initial_state: torch.Tensor = None,
119
output_final_state: bool = False,
120
cu_seqlens: torch.LongTensor | None = None,
122
- head_first: bool = False,
121
use_qk_l2norm_in_kernel: bool = False,
122
):
123
r"""
124
Args:
125
q (torch.Tensor):
128
- queries of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`.
126
+ Queries of shape `[B, T, H, K]`.
127
k (torch.Tensor):
130
- keys of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`.
128
+ Keys of shape `[B, T, H, K]`.
129
v (torch.Tensor):
132
- values of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`.
130
+ Values of shape `[B, T, H, V]`.
131
g (torch.Tensor):
134
- (forget) gating tensor (in log space!) of shape `[B, T, H]` if `head_first=False` else `[B, H, T]`.
132
+ (forget) Gating tensor (in log space!) of shape `[B, T, H]`.
133
beta (torch.Tensor):
136
- betas of shape `[B, T, H]` if `head_first=False` else `[B, H, T]`.
134
+ Betas of shape `[B, T, H]`.
135
scale (Optional[int]):
136
Scale factor for the RetNet attention scores.
137
If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
@@ -146,13 +144,9 @@ def chunk_gated_delta_rule(
144
cu_seqlens (torch.LongTensor):
145
Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
146
consistent with the FlashAttention API.
149
- head_first (Optional[bool]):
150
- Whether the inputs are in the head-first format, which is not supported for variable-length inputs.
151
- Default: `False`.
152
-
147
Returns:
148
o (torch.Tensor):
155
- Outputs of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`.
149
+ Outputs of shape `[B, T, H, V]`.
150
final_state (torch.Tensor):
151
Final state of shape `[N, H, V, K]` if `output_final_state=True` else `None`.
152
@@ -189,24 +183,11 @@ def chunk_gated_delta_rule(
183
assert q.dtype != torch.float32, (
184
"ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16."
185
)
192
- assert len(beta.shape) == 3, (
193
- "beta must be of shape [B, T, H] if head_first=False, or [B, H, T] otherwise."
194
- )
195
-
196
- if head_first:
197
- raise DeprecationWarning(
198
- "head_first is deprecated and will be removed in a future version. "
199
- "Please use head_first=False for now instead.",
200
- stacklevel=2,
201
- )
202
- q, k, v, beta, g = map(
203
- lambda x: rearrange(x, "b h t ... -> b t h ..."), (q, k, v, beta, g)
204
- )
205
- if not head_first and q.shape[1] < q.shape[2]:
186
+ assert len(beta.shape) == 3, "beta must be of shape [B, T, H]."
187
+ if q.shape[1] < q.shape[2]:
188
warnings.warn(
189
f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). "
190
"This may indicate the inputs were passed in head-first format [B, H, T, ...] "
209
- "when head_first=False was specified. "
191
"Please verify your input tensor format matches the expected shape [B, T, H, ...].",
192
stacklevel=2,
193
)
@@ -235,6 +216,4 @@ def chunk_gated_delta_rule(
216
cu_seqlens,
217
use_qk_l2norm_in_kernel,
218
)
238
- if head_first:
239
- o = rearrange(o, "b t h ... -> b h t ...")
219
return o, final_state
vllm/model_executor/models/llava_onevision.py
-1
@@ -867,7 +867,6 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
867
mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
868
if not mm_input_by_modality:
869
return []
870
- return None
870
871
# The result multimodal_embeddings is tuple of tensors, with each
872
# tensor corresponding to a multimodal data item (image or video).
vllm/model_executor/models/qwen3_next.py
-6
@@ -115,7 +115,6 @@ def fi_chunk_gated_delta_rule(
115
initial_state: torch.Tensor,
116
output_final_state: bool,
117
cu_seqlens: torch.LongTensor | None = None,
118
- head_first: bool = False,
118
use_qk_l2norm_in_kernel: bool = True,
119
):
120
from flashinfer.gdn_prefill import (
@@ -172,7 +171,6 @@ class ChunkGatedDeltaRule(CustomOp):
171
initial_state: torch.Tensor,
172
output_final_state: bool,
173
cu_seqlens: torch.LongTensor | None = None,
175
- head_first: bool = False,
174
use_qk_l2norm_in_kernel: bool = True,
175
):
176
return fi_chunk_gated_delta_rule(
@@ -184,7 +182,6 @@ class ChunkGatedDeltaRule(CustomOp):
182
initial_state=initial_state,
183
output_final_state=output_final_state,
184
cu_seqlens=cu_seqlens,
187
- head_first=head_first,
185
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
186
)
187
@@ -198,7 +195,6 @@ class ChunkGatedDeltaRule(CustomOp):
195
initial_state: torch.Tensor,
196
output_final_state: bool,
197
cu_seqlens: torch.LongTensor | None = None,
201
- head_first: bool = False,
198
use_qk_l2norm_in_kernel: bool = True,
199
):
200
return fla_chunk_gated_delta_rule(
@@ -210,7 +206,6 @@ class ChunkGatedDeltaRule(CustomOp):
206
initial_state=initial_state,
207
output_final_state=output_final_state,
208
cu_seqlens=cu_seqlens,
213
- head_first=head_first,
209
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
210
)
211
@@ -790,7 +785,6 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
785
initial_state=initial_state,
786
output_final_state=True,
787
cu_seqlens=non_spec_query_start_loc,
793
- head_first=False,
788
use_qk_l2norm_in_kernel=True,
789
)
790
# Init cache