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 5, 19:45 UTC
[Bugfix] Remove false-positive format mismatch warnings in FLA ops (#38255)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
Thomas Parnell committed
Mar 30, 2026 at 14:32 UTC
7c3f88b2a895f58d653e9284dbb7b1da85c85e73
2 files changed
-17
vllm/model_executor/layers/fla/ops/chunk.py
-8
@@ -7,7 +7,6 @@
7
# the following copyright notice:
8
# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
9
# ruff: noqa: E501
10
-import warnings
10
11
import torch
12
@@ -184,13 +183,6 @@ def chunk_gated_delta_rule(
183
"ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16."
184
)
185
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, ...] "
191
- "Please verify your input tensor format matches the expected shape [B, T, H, ...].",
192
- stacklevel=2,
193
- )
186
if cu_seqlens is not None:
187
if q.shape[0] != 1:
188
raise ValueError(
vllm/model_executor/layers/fla/ops/cumsum.py
-9
@@ -7,7 +7,6 @@
7
# the following copyright notice:
8
# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
9
# ruff: noqa: E501
10
-import warnings
10
11
import torch
12
@@ -252,14 +251,6 @@ def chunk_local_cumsum(
251
output_dtype: torch.dtype | None = torch.float,
252
**kwargs,
253
) -> torch.Tensor:
255
- if not head_first and g.shape[1] < g.shape[2]:
256
- warnings.warn(
257
- f"Input tensor shape suggests potential format mismatch: seq_len ({g.shape[1]}) < num_heads ({g.shape[2]}). "
258
- "This may indicate the inputs were passed in head-first format [B, H, T, ...] "
259
- "when head_first=False was specified. "
260
- "Please verify your input tensor format matches the expected shape [B, T, H, ...].",
261
- stacklevel=2,
262
- )
254
if cu_seqlens is not None:
255
assert g.shape[0] == 1, (
256
"Only batch size 1 is supported when cu_seqlens are provided"