[PCG] fix piecewise cuda graph for Qwen3.5 (#19220)

This commit is contained in:
Minglei Zhu
2026-02-25 19:16:52 -08:00
committed by GitHub
parent a0a8f1473c
commit b3202fe6d0
4 changed files with 9 additions and 46 deletions

View File

@@ -72,6 +72,13 @@ if _is_cuda:
N = mat_b.shape[-1]
return mat_a.new_empty((M, N), dtype=out_dtype)
@torch.library.register_fake("sgl_kernel::fp8_blockwise_scaled_mm")
def _fp8_blockwise_scaled_mm_abstract(mat_a, mat_b, scales_a, scales_b, out_dtype):
# mat_a: [M, K], mat_b: [K, N] or [N, K] depending on callsite layout; output is [M, N].
M = mat_a.shape[-2]
N = mat_b.shape[-1]
return mat_a.new_empty((M, N), dtype=out_dtype)
use_vllm_cutlass_w8a8_fp8_kernel = get_bool_env_var("USE_VLLM_CUTLASS_W8A8_FP8_KERNEL")
use_triton_w8a8_fp8_kernel = get_bool_env_var("USE_TRITON_W8A8_FP8_KERNEL")