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, 04:45 UTC
[Bugfix] Fix mismatch between global and local attention heads in tensor-parallel mode for param2moe model (#39707)
Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
bhargav-patel-29 committed
Apr 14, 2026 at 17:43 UTC
f7e62e3d6618f64430262bf776079d0d89f20501
1 file changed
+5
-11
vllm/model_executor/models/param2moe.py
+5
-11
@@ -16,7 +16,6 @@
16
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17
# See the License for the specific language governing permissions and
18
# limitations under the License.
19
-# limitations under the License.
19
from __future__ import annotations
20
21
from collections.abc import Iterable, Iterator
@@ -202,10 +201,10 @@ class Param2MoEAttention(nn.Module):
201
)
202
203
self.attn = Attention(
205
- num_heads=self.num_heads,
204
+ num_heads=self.num_local_heads,
205
head_size=self.head_dim,
206
scale=self.scaling,
208
- num_kv_heads=self.num_kv_heads,
207
+ num_kv_heads=self.num_local_kv_heads,
208
cache_config=cache_config,
209
quant_config=quant_config,
210
prefix=f"{prefix}.attn",
@@ -216,15 +215,15 @@ class Param2MoEAttention(nn.Module):
215
positions: torch.Tensor,
216
hidden_states: torch.Tensor,
217
) -> torch.Tensor:
219
- # 1. Fused QKV projection → split into local Q / K / V
218
qkv, _ = self.qkv_proj(hidden_states)
219
q, k, v = qkv.split(
220
[self.q_size_local, self.kv_size_local, self.kv_size_local],
221
dim=-1,
222
)
223
+ q = q.contiguous()
224
+ k = k.contiguous()
225
+ v = v.contiguous()
226
226
- # 2. Optional per-head QK norms
227
- # Reshape to (T, num_local_heads, head_dim), norm, reshape back.
227
if self.use_qk_norm:
228
T = q.shape[0]
229
q = self.q_layernorm(q.view(T, self.num_local_heads, self.head_dim)).view(
@@ -234,13 +233,8 @@ class Param2MoEAttention(nn.Module):
233
k.view(T, self.num_local_kv_heads, self.head_dim)
234
).view(T, self.kv_size_local)
235
237
- # 3. Rotary position embeddings
236
q, k = self.rotary_emb(positions, q, k)
239
-
240
- # 4. Paged attention
237
attn_output = self.attn(q, k, v)
242
-
243
- # 5. Output projection
238
output, _ = self.o_proj(attn_output)
239
return output
240