A high-throughput and memory-efficient inference and serving engine for LLMs
Revert "[Kernel] Add PDL support for DeepGEMM kernel" (#45999)
Micah Williamson committed
Jun 18, 2026 at 00:59 UTC
e945169207ac90e0da4f21f579f309b28caabe90
3 files changed
+20
-62
csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu
+14
-35
@@ -304,17 +304,9 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
304
const int sf_k_idx = blockIdx.x * kGroupsPerBlockX + sf_k_local;
305
const int mn_idx = blockIdx.y * kRowsPerBlock + row_local;
306
307
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
308
- asm volatile("griddepcontrol.wait;");
309
-#endif
310
-
307
if (mn_idx >= tma_aligned_mn) {
312
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
313
- asm volatile("griddepcontrol.launch_dependents;");
314
-#endif
308
return;
309
}
317
-
310
const bool is_valid_group = (mn_idx < mn) && (sf_k_idx < groups_per_row);
311
312
// Load 16 input elements (32 B) into registers as two adjacent uint4
@@ -425,10 +417,6 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
417
static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
418
sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE;
419
*reinterpret_cast<uint4*>(group_output) = packed_out;
428
-
429
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
430
- asm volatile("griddepcontrol.launch_dependents;");
431
-#endif
420
}
421
422
// Public entry point: register-resident packed quant kernel.
@@ -509,29 +497,20 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
497
498
#define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \
499
do { \
512
- cudaLaunchConfig_t config = {}; \
513
- config.gridDim = dim3(static_cast<unsigned int>(blocks_x), \
514
- static_cast<unsigned int>(blocks_y)); \
515
- config.blockDim = dim3(num_threads); \
516
- config.dynamicSmemBytes = 0; \
517
- config.stream = stream; \
518
- cudaLaunchAttribute attrs[1]; \
519
- attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; \
520
- attrs[0].val.programmaticStreamSerializationAllowed = 1; \
521
- config.numAttrs = 1; \
522
- config.attrs = attrs; \
523
- cudaLaunchKernelEx( \
524
- &config, \
525
- per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, \
526
- KX, RY>, \
527
- static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
528
- reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
529
- static_cast<int>(padded_groups_per_row), \
530
- static_cast<int>(groups_per_row), static_cast<int>(mn), \
531
- static_cast<int>(output_q_mn_extent), \
532
- static_cast<int>(tma_aligned_mn), num_scale_elems, \
533
- static_cast<float>(eps), static_cast<float>(min_8bit), \
534
- static_cast<float>(max_8bit)); \
500
+ dim3 grid(static_cast<unsigned int>(blocks_x), \
501
+ static_cast<unsigned int>(blocks_y)); \
502
+ dim3 block(num_threads); \
503
+ per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, KX, \
504
+ RY> \
505
+ <<<grid, block, 0, stream>>>( \
506
+ static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
507
+ reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
508
+ static_cast<int>(padded_groups_per_row), \
509
+ static_cast<int>(groups_per_row), static_cast<int>(mn), \
510
+ static_cast<int>(output_q_mn_extent), \
511
+ static_cast<int>(tma_aligned_mn), num_scale_elems, \
512
+ static_cast<float>(eps), static_cast<float>(min_8bit), \
513
+ static_cast<float>(max_8bit)); \
514
} while (0)
515
516
#define LAUNCH_REG_KERNEL(T, DST_DTYPE) \
vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py
+6
-8
@@ -37,8 +37,6 @@ def _fused_inv_rope_fp8_quant_per_head(
37
ROPE_START: tl.constexpr,
38
HALF_ROPE: tl.constexpr,
39
TMA_ALIGNED_SCALES: tl.constexpr,
40
- USE_GDC: tl.constexpr,
41
- launch_pdl: tl.constexpr, # triton metadata
40
):
41
# int64: stride multiply overflows int32 past num_tokens=32768 (IMA).
42
pid_token = tl.program_id(0).to(tl.int64)
@@ -48,9 +46,7 @@ def _fused_inv_rope_fp8_quant_per_head(
46
head_in_group = pid_gh % heads_per_group
47
global_head = pid_gh
48
qb_start = head_in_group * CHUNKS_PER_HEAD
51
- if USE_GDC:
52
- tl.extra.cuda.gdc_launch_dependents()
53
- tl.extra.cuda.gdc_wait()
49
+
50
# Padding rows in the TMA-aligned scale buffer: fill with zero and skip quant.
51
if pid_token >= num_tokens:
52
if TMA_ALIGNED_SCALES:
@@ -247,8 +243,11 @@ def _fused_inv_rope_fp8_quant_kernel_impl(
243
(scale_inner * tma_aligned_T, 1, tma_aligned_T),
244
)
245
grid = (tma_aligned_T, n_groups * heads_per_group)
250
- use_gdc = current_platform.is_arch_support_pdl()
251
- pdl_kwargs = {"launch_pdl": True} if use_gdc else {}
246
+ pdl_kwargs = (
247
+ {}
248
+ if current_platform.is_rocm() or current_platform.is_xpu()
249
+ else {"launch_pdl": False}
250
+ )
251
_fused_inv_rope_fp8_quant_per_head[grid](
252
o,
253
positions,
@@ -271,7 +270,6 @@ def _fused_inv_rope_fp8_quant_kernel_impl(
270
ROPE_START=rope_start,
271
HALF_ROPE=half_rope,
272
TMA_ALIGNED_SCALES=tma_aligned_scales,
274
- USE_GDC=use_gdc,
273
num_stages=1,
274
**pdl_kwargs,
275
num_warps=1,
vllm/utils/deep_gemm.py
-19
@@ -177,22 +177,6 @@ def _import_deep_gemm():
177
return None
178
179
180
-def _apply_pdl(mod, enable: bool = True) -> None:
181
- mod_name = getattr(mod, "__name__", str(mod))
182
- try:
183
- set_pdl_fn = getattr(mod, "set_pdl", None)
184
- if set_pdl_fn is None:
185
- return
186
- set_pdl_fn(enable)
187
- logger.info_once(
188
- "DeepGEMM PDL %s on %s.",
189
- "enabled" if enable else "disabled",
190
- mod_name,
191
- )
192
- except Exception as e: # noqa: BLE001
193
- logger.warning_once("Failed to set DeepGEMM PDL on %s: %s", mod_name, e)
194
-
195
-
180
def _lazy_init() -> None:
181
"""Import deep_gemm and resolve symbols on first use."""
182
global _cublaslt_gemm_nt_impl
@@ -235,9 +219,6 @@ def _lazy_init() -> None:
219
if _dg is None:
220
return
221
238
- # Enable PDL for DeepGEMM on architectures that support it (SM90+).
239
- if current_platform.is_arch_support_pdl():
240
- _apply_pdl(_dg, True)
222
_cublaslt_gemm_nt_impl = getattr(_dg, "cublaslt_gemm_nt", None)
223
_fp8_gemm_nt_impl = getattr(_dg, "fp8_gemm_nt", None)
224
_fp8_einsum_impl = getattr(_dg, "fp8_einsum", None)