[Performance][DSv4] Size dequant gather launch grid by rows (#55061)
Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: aoshen02 <aoshen524@gmail.com>
aoshen02 committed
Sep 4, 2026 at 11:06 UTC
69cf055936066779cbd54c237739e9ac657ed86b
1 file changed
+6
-1
vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py
+6
-1
@@ -73,7 +73,12 @@ class DequantGatherKCacheKernel:
73
),
74
)
75
76
- grid = (out.shape[0], 1024, 1)
76
+ num_reqs = out.shape[0]
77
+ max_rows = out.shape[1] - offset
78
+ num_workers = cutlass.max(
79
+ 1, cutlass.min(cute.ceil_div(max_rows, 4), cute.ceil_div(8192, num_reqs))
80
+ )
81
+ grid = (num_reqs, num_workers, 1)
82
self.kernel(
83
out,
84
k_data,