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 11, 13:45 UTC
[Perf] Set split_k to 1 for triton_kernels (#30528)
Signed-off-by: Xin Yang <xyangx@amazon.com>
Xin Yang committed
Dec 12, 2025 at 11:07 UTC
1f19d8f899b228a530d256bf9476d9b1ea3039af
1 file changed
+12
-6
vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
+12
-6
@@ -57,12 +57,18 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps):
57
mx_axis=1, num_warps=num_warps
58
)
59
)
60
- if current_platform.is_cuda() and current_platform.is_device_capability(100):
61
- constraints = {
62
- "is_persistent": True,
63
- "epilogue_subtile": 1,
64
- }
65
- opt_flags.update_opt_flags_constraints(constraints)
60
+ if current_platform.is_cuda():
61
+ if current_platform.is_device_capability(90):
62
+ constraints = {
63
+ "split_k": 1,
64
+ }
65
+ opt_flags.update_opt_flags_constraints(constraints)
66
+ elif current_platform.is_device_capability(100):
67
+ constraints = {
68
+ "is_persistent": True,
69
+ "epilogue_subtile": 1,
70
+ }
71
+ opt_flags.update_opt_flags_constraints(constraints)
72
# transpose the tensor so that the quantization axis is on dim1
73
quant_tensor = quant_tensor.transpose(-2, -1)
74
scale = scale.transpose(-2, -1)