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, 07:30 UTC
Revert "[Bugfix] Fix fused MoE IMA (sans chunking) by using int64 for strides" (#34530)
Michael Goin committed
Feb 13, 2026 at 13:35 UTC
bfaa5593050ec9bf60e2361b3b9dc575efeee83f
1 file changed
+27
-27
vllm/model_executor/layers/fused_moe/fused_moe.py
+27
-27
@@ -98,19 +98,19 @@ def fused_moe_kernel_gptq_awq(
98
# moving by 1 element in a particular dimension. E.g. `stride_am` is
99
# how much to increase `a_ptr` by to get the element one row down
100
# (A has M rows).
101
- stride_am: tl.int64,
102
- stride_ak: tl.int64,
103
- stride_be: tl.int64,
104
- stride_bk: tl.int64,
105
- stride_bn: tl.int64,
106
- stride_cm: tl.int64,
107
- stride_cn: tl.int64,
108
- stride_bse: tl.int64,
109
- stride_bsk: tl.int64,
110
- stride_bsn: tl.int64,
111
- stride_bze: tl.int64,
112
- stride_bzk: tl.int64,
113
- stride_bzn: tl.int64,
101
+ stride_am,
102
+ stride_ak,
103
+ stride_be,
104
+ stride_bk,
105
+ stride_bn,
106
+ stride_cm,
107
+ stride_cn,
108
+ stride_bse,
109
+ stride_bsk,
110
+ stride_bsn,
111
+ stride_bze,
112
+ stride_bzk,
113
+ stride_bzn,
114
block_k_diviable: tl.constexpr,
115
group_size: tl.constexpr,
116
# Meta-parameters
@@ -332,20 +332,20 @@ def fused_moe_kernel(
332
# moving by 1 element in a particular dimension. E.g. `stride_am` is
333
# how much to increase `a_ptr` by to get the element one row down
334
# (A has M rows).
335
- stride_am: tl.int64,
336
- stride_ak: tl.int64,
337
- stride_be: tl.int64,
338
- stride_bk: tl.int64,
339
- stride_bn: tl.int64,
340
- stride_cm: tl.int64,
341
- stride_cn: tl.int64,
342
- stride_asm: tl.int64,
343
- stride_ask: tl.int64,
344
- stride_bse: tl.int64,
345
- stride_bsk: tl.int64,
346
- stride_bsn: tl.int64,
347
- stride_bbe: tl.int64, # bias expert stride
348
- stride_bbn: tl.int64, # bias N stride
335
+ stride_am,
336
+ stride_ak,
337
+ stride_be,
338
+ stride_bk,
339
+ stride_bn,
340
+ stride_cm,
341
+ stride_cn,
342
+ stride_asm,
343
+ stride_ask,
344
+ stride_bse,
345
+ stride_bsk,
346
+ stride_bsn,
347
+ stride_bbe, # bias expert stride
348
+ stride_bbn, # bias N stride
349
# Block size for block-wise quantization
350
group_n: tl.constexpr,
351
group_k: tl.constexpr,