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, 12:30 UTC
[Bugfix] Fix fused MoE IMA (sans chunking) by using int64 for strides (#34279)
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Tyler Michael Smith committed
Feb 11, 2026 at 00:15 UTC
d7982daff5334b9465b29fa943a1954c064ab226
1 file changed
+27
-27
vllm/model_executor/layers/fused_moe/fused_moe.py
+27
-27
@@ -95,19 +95,19 @@ def fused_moe_kernel_gptq_awq(
95
# moving by 1 element in a particular dimension. E.g. `stride_am` is
96
# how much to increase `a_ptr` by to get the element one row down
97
# (A has M rows).
98
- stride_am,
99
- stride_ak,
100
- stride_be,
101
- stride_bk,
102
- stride_bn,
103
- stride_cm,
104
- stride_cn,
105
- stride_bse,
106
- stride_bsk,
107
- stride_bsn,
108
- stride_bze,
109
- stride_bzk,
110
- stride_bzn,
98
+ stride_am: tl.int64,
99
+ stride_ak: tl.int64,
100
+ stride_be: tl.int64,
101
+ stride_bk: tl.int64,
102
+ stride_bn: tl.int64,
103
+ stride_cm: tl.int64,
104
+ stride_cn: tl.int64,
105
+ stride_bse: tl.int64,
106
+ stride_bsk: tl.int64,
107
+ stride_bsn: tl.int64,
108
+ stride_bze: tl.int64,
109
+ stride_bzk: tl.int64,
110
+ stride_bzn: tl.int64,
111
block_k_diviable: tl.constexpr,
112
group_size: tl.constexpr,
113
# Meta-parameters
@@ -329,20 +329,20 @@ def fused_moe_kernel(
329
# moving by 1 element in a particular dimension. E.g. `stride_am` is
330
# how much to increase `a_ptr` by to get the element one row down
331
# (A has M rows).
332
- stride_am,
333
- stride_ak,
334
- stride_be,
335
- stride_bk,
336
- stride_bn,
337
- stride_cm,
338
- stride_cn,
339
- stride_asm,
340
- stride_ask,
341
- stride_bse,
342
- stride_bsk,
343
- stride_bsn,
344
- stride_bbe, # bias expert stride
345
- stride_bbn, # bias N stride
332
+ stride_am: tl.int64,
333
+ stride_ak: tl.int64,
334
+ stride_be: tl.int64,
335
+ stride_bk: tl.int64,
336
+ stride_bn: tl.int64,
337
+ stride_cm: tl.int64,
338
+ stride_cn: tl.int64,
339
+ stride_asm: tl.int64,
340
+ stride_ask: tl.int64,
341
+ stride_bse: tl.int64,
342
+ stride_bsk: tl.int64,
343
+ stride_bsn: tl.int64,
344
+ stride_bbe: tl.int64, # bias expert stride
345
+ stride_bbn: tl.int64, # bias N stride
346
# Block size for block-wise quantization
347
group_n: tl.constexpr,
348
group_k: tl.constexpr,