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 15, 17:15 UTC
Use Cache Hinting for fused_moe kernel (#15511)
Wes committed
Mar 26, 2025 at 17:21 UTC
7a888271f5bd401f8fc64704c239833244471a91
1 file changed
+12
-4
vllm/model_executor/layers/fused_moe/fused_moe.py
+12
-4
@@ -189,7 +189,11 @@ def fused_moe_kernel_gptq_awq(
189
mask=token_mask[:, None] &
190
(offs_k[None, :] < K - k * BLOCK_SIZE_K),
191
other=0.0)
192
- b = tl.load(b_ptrs)
192
+ b = tl.load(
193
+ b_ptrs,
194
+ cache_modifier=".cg",
195
+ eviction_policy="evict_last",
196
+ )
197
if use_int4_w4a16:
198
b = (b >> b_shifter) & 0xF
199
@@ -391,9 +395,13 @@ def fused_moe_kernel(
395
mask=token_mask[:, None] &
396
(offs_k[None, :] < K - k * BLOCK_SIZE_K),
397
other=0.0)
394
- b = tl.load(b_ptrs,
395
- mask=offs_k[:, None] < K - k * BLOCK_SIZE_K,
396
- other=0.0)
398
+ b = tl.load(
399
+ b_ptrs,
400
+ mask=offs_k[:, None] < K - k * BLOCK_SIZE_K,
401
+ other=0.0,
402
+ cache_modifier=".cg",
403
+ eviction_policy="evict_last",
404
+ )
405
# We accumulate along the K dimension.
406
if use_int8_w8a16:
407
accumulator = tl.dot(a, b.to(compute_type), acc=accumulator)