Merge branch 'main' into patches

Signed-off-by: MOVZX <movzx@yahoo.com>

csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_bfloat16_fe2m1f_bfloat16.cu

@@ -0,0 +1,70 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 256, 1, 8, 8, true, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 8, 4, true, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 4, 8, true, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 256, 1, 8, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 8, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 4, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 256, 2, 16, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 2, 8, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 2, 4, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 256, 3, 16, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 3, 8, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 3, 4, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 256, 4, 16, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 4, 8, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 4, 4, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 256, 1, 8, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 1, 8, 4, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 1, 4, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_bfloat16_fe4m3fn_bfloat16.cu

@@ -0,0 +1,100 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 256, 1, 8, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 1, 8, 4, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 1, 4, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_bfloat16_u4_bfloat16.cu

@@ -0,0 +1,130 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_bfloat16_u4b8_bfloat16.cu

@@ -0,0 +1,130 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_bfloat16_u8b128_bfloat16.cu

@@ -0,0 +1,130 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_float16_fe2m1f_float16.cu

@@ -0,0 +1,40 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 256, 1, 8, 8, true, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 8, 4, true, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 4, 8, true, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 256, 1, 8, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 8, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 4, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 256, 2, 16, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 128, 2, 8, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 128, 2, 4, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 256, 3, 16, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 128, 3, 8, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 128, 3, 4, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 256, 4, 16, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 128, 4, 8, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 128, 4, 4, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_float16_fe4m3fn_float16.cu

@@ -0,0 +1,70 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_float16_u4_float16.cu

@@ -0,0 +1,130 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_float16_u4b8_float16.cu

@@ -0,0 +1,130 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_float16_u8b128_float16.cu

@@ -0,0 +1,130 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_s8_u4_bfloat16.cu

@@ -0,0 +1,106 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_s8_u4_float16.cu

@@ -0,0 +1,106 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_s8_u4b8_bfloat16.cu

@@ -0,0 +1,106 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_s8_u4b8_float16.cu

@@ -0,0 +1,106 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/quantization/marlin/sm80_kernel_bfloat16_fe2m1f_bfloat16.cu

@@ -0,0 +1,70 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 256, 1, 8, 8, true, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 8, 4, true, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 4, 8, true, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 256, 1, 8, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 8, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 4, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 256, 2, 16, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 2, 8, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 2, 4, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 256, 3, 16, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 3, 8, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 3, 4, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 256, 4, 16, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 4, 8, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 4, 4, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 256, 1, 8, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 1, 8, 4, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 1, 4, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/quantization/marlin/sm80_kernel_bfloat16_fe4m3fn_bfloat16.cu

@@ -0,0 +1,100 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 256, 1, 8, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 1, 8, 4, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 1, 4, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), vllm::kBFloat16.id(), vllm::kFE8M0fnu.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/quantization/marlin/sm80_kernel_bfloat16_u4_bfloat16.cu

@@ -0,0 +1,130 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/quantization/marlin/sm80_kernel_bfloat16_u4b8_bfloat16.cu

@@ -0,0 +1,130 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/quantization/marlin/sm80_kernel_bfloat16_u8b128_bfloat16.cu

@@ -0,0 +1,130 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU8B128.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/quantization/marlin/sm80_kernel_float16_fe2m1f_float16.cu

@@ -0,0 +1,40 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 256, 1, 8, 8, true, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 8, 4, true, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 4, 8, true, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 256, 1, 8, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 8, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 4, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 256, 2, 16, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 128, 2, 8, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 128, 2, 4, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 256, 3, 16, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 128, 3, 8, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 128, 3, 4, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 256, 4, 16, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 128, 4, 8, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE2M1f.id(), vllm::kFloat16.id(), vllm::kFE4M3fn.id(), 128, 4, 4, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/quantization/marlin/sm80_kernel_float16_fe4m3fn_float16.cu

@@ -0,0 +1,70 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kFE4M3fn.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/quantization/marlin/sm80_kernel_float16_u4_float16.cu

@@ -0,0 +1,130 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/quantization/marlin/sm80_kernel_float16_u4b8_float16.cu

@@ -0,0 +1,130 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/quantization/marlin/sm80_kernel_float16_u8b128_float16.cu

@@ -0,0 +1,130 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU8B128.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/quantization/marlin/sm80_kernel_s8_u4_bfloat16.cu

@@ -0,0 +1,106 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/quantization/marlin/sm80_kernel_s8_u4_float16.cu

@@ -0,0 +1,106 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/quantization/marlin/sm80_kernel_s8_u4b8_bfloat16.cu

@@ -0,0 +1,106 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

csrc/libtorch_stable/quantization/marlin/sm80_kernel_s8_u4b8_float16.cu

@@ -0,0 +1,106 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kS8.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}

docs/features/engram.md

@@ -83,7 +83,7 @@ directly from the Triton kernel through unified virtual addressing (UVA),
 prefetching on a side CUDA stream so transfers overlap with the compute of
 preceding layers. The freed GPU memory goes to the KV cache.
 
-`cpu_offload` defaults to `true` (on)
+`cpu_offload` defaults to `true` and is set only through `--engram-config`.
 
 To keep the tables resident on GPU instead (e.g. when host memory is
 scarce):
@@ -93,8 +93,40 @@ vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
   --engram-config.cpu_offload false
 ```
 
-CPU offload requires a GPU with UVA support;
+CPU offload requires a GPU with UVA support; vLLM fails fast when it is
-vLLM fails fast when it is unavailable.
+unavailable.
+
+## Host file gather (Qwen4Exp)
+
+`host_file_gather` serves Qwen4Exp PLE rows from the loader's own file
+mapping of the safetensors shards instead of allocating the table. Set
+`--engram-config '{"host_file_gather": true}'`. It applies only to Qwen4Exp
+and takes precedence over `cpu_offload`.
+
+Rows stay in the checkpoint dtype: an FP8 table stays FP8 and is dequantized
+on the GPU, and an unquantized table must match `--dtype`. CPU offload
+converts on load; this mode rejects the mismatch at load.
+
+Faulted pages are reclaimable page cache, shared across co-located processes.
+They appear in each process's RSS as file-backed (`RssFile`), and PSS is
+where they are split across processes. A write to the mapping copies that
+page into anonymous RSS. Each step pays a host sync, and every TP rank
+gathers all rows. Co-located ranks share the cache; ranks on different
+machines do not.
+
+Serving depends on the checkpoint files staying unchanged while the server
+runs: an in-place overwrite serves the new bytes, truncation kills the worker
+with SIGBUS, and deleting or renaming the files is safe.
+
+The first touch of a row costs a disk read, and a page evicted under memory
+pressure is read again on the next miss. By default no prewarm is needed; for
+a fast first request, read the PLE files once after the server has started
+and allocated its KV cache.
+
+Reload is unsupported and fails closed. ROCm and elastic EP reject the key.
+Load paths that do not hand over file-backed safetensors tensors fail at bind
+(`eager`, `.bin`, runai, `fastsafetensors`, `instanttensor`). The multithread
+loader is file-backed and is admitted. `--load-format dummy` serves zeros.
 
 ## Data-parallel topologies (DeepSeek V4.1)
 
@@ -115,7 +147,8 @@ vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
 ```
 
 `embedding_across_dp` is not supported with elastic expert parallelism yet.
-Qwen4Exp PLE tables are ETP-sharded instead and honor `cpu_offload` only.
+Qwen4Exp PLE tables are ETP-sharded instead and honor only `cpu_offload` and
+`host_file_gather`.
 
 ## Sharing host tables across DP replicas
 

tests/config/test_recoverssm_config.py

@@ -0,0 +1,172 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Configuration-level regression tests for Qwen4Exp GDN/PLE RecoverSSM selection.
+
+``--use-replayssm`` with speculative decoding selects RecoverSSM for Qwen4Exp, as it
+does for Kimi-K3 KDA. The MTP drafter's derived config shares ``cache_config`` and must
+keep the choice. These tests run the real validator function against stand-in configs.
+"""
+
+from types import SimpleNamespace
+from typing import Any, cast
+
+import pytest
+
+from vllm.config.mamba import MambaBackendEnum
+from vllm.config.vllm import VllmConfig
+from vllm.model_executor.layers.mamba.recoverssm_utils import uses_recoverssm
+from vllm.v1.attention.backend import AttentionCGSupport
+
+TARGET = "Qwen4ExpForConditionalGeneration"
+DRAFTER = "Qwen4ExpMTP"
+
+
+def _validator():
+    validators = VllmConfig.__pydantic_decorators__.model_validators  # type: ignore[attr-defined]
+    return validators["validate_mamba_cached_kernel"].func
+
+
+def _config(
+    arch: str,
+    cache_config=None,
+    num_spec: int = 3,
+    use_replayssm: bool = True,
+    supports_replayssm: bool = True,
+    **overrides,
+):
+    cfg = SimpleNamespace(
+        num_speculative_tokens=num_spec,
+        model_config=SimpleNamespace(
+            architecture=arch, supports_replayssm=supports_replayssm
+        ),
+        cache_config=cache_config
+        or SimpleNamespace(
+            use_recoverssm=False,
+            use_replayssm=use_replayssm,
+            mamba_cache_mode="align",
+        ),
+        mamba_config=SimpleNamespace(
+            enable_stochastic_rounding=False, backend=MambaBackendEnum.TRITON
+        ),
+        use_v2_model_runner=True,
+        parallel_config=SimpleNamespace(pipeline_parallel_size=1),
+        kv_transfer_config=None,
+    )
+    for key, value in overrides.items():
+        target, _, attr = key.partition("__")
+        if attr:
+            setattr(getattr(cfg, target), attr, value)
+        else:
+            setattr(cfg, target, value)
+    return cfg
+
+
+def test_validator_is_registered():
+    assert "validate_mamba_cached_kernel" in (
+        VllmConfig.__pydantic_decorators__.model_validators  # type: ignore[attr-defined]
+    )
+
+
+@pytest.mark.parametrize(
+    "arch", ["Qwen4ExpForConditionalGeneration", "Qwen4ExpForCausalLM"]
+)
+def test_replayssm_with_spec_selects_recoverssm(arch):
+    cfg = _config(arch)
+    _validator()(cfg)
+    assert cfg.cache_config.use_recoverssm is True
+    assert uses_recoverssm(cfg.cache_config, 3)
+
+
+def test_derived_drafter_keeps_the_choice_and_runs_the_checks():
+    target = _config(TARGET)
+    _validator()(target)
+    drafter = _config(DRAFTER, cache_config=target.cache_config)
+    _validator()(drafter)
+    assert target.cache_config.use_recoverssm is True
+    bad = _config(
+        DRAFTER,
+        cache_config=target.cache_config,
+        parallel_config__pipeline_parallel_size=2,
+    )
+    with pytest.raises(ValueError, match="pipeline_parallel_size"):
+        _validator()(bad)
+
+
+@pytest.mark.parametrize(
+    ("override", "value", "match"),
+    [
+        ("parallel_config__pipeline_parallel_size", 2, "pipeline_parallel_size"),
+        ("mamba_config__enable_stochastic_rounding", True, "stochastic"),
+        ("mamba_config__backend", "flashinfer", "mamba-backend triton"),
+        ("use_v2_model_runner", False, "V2_MODEL_RUNNER"),
+    ],
+)
+def test_runtime_checks_reject(override, value, match):
+    cfg = _config(TARGET, **{override: value})
+    with pytest.raises(ValueError, match=match):
+        _validator()(cfg)
+
+
+def test_flag_unset_keeps_the_stock_path():
+    cfg = _config(TARGET, use_replayssm=False)
+    _validator()(cfg)
+    assert cfg.cache_config.use_recoverssm is False
+    assert not uses_recoverssm(cfg.cache_config, 3)
+
+
+def test_qwen4exp_replayssm_requires_speculative_decoding():
+    with pytest.raises(ValueError, match="requires speculative decoding"):
+        _validator()(_config(TARGET, num_spec=0))
+
+
+def test_architecture_without_replayssm_support_is_rejected():
+    cfg = _config("Qwen3NextForCausalLM", supports_replayssm=False)
+    with pytest.raises(ValueError, match="not supported for architecture"):
+        _validator()(cfg)
+
+
+def test_uses_recoverssm_predicate():
+    on = SimpleNamespace(use_recoverssm=True)
+    assert uses_recoverssm(on, 3)
+    assert not uses_recoverssm(on, 0)
+    assert not uses_recoverssm(SimpleNamespace(use_recoverssm=False), 3)
+    assert not uses_recoverssm(SimpleNamespace(), 3)
+
+
+def test_builders_fall_back_to_piecewise():
+    from vllm.v1.attention.backends.gdn_recoverssm import GDNRecoverSSMMetadataBuilder
+    from vllm.v1.attention.backends.ple_recoverssm import PleRecoverSSMMetadataBuilder
+
+    for builder in (GDNRecoverSSMMetadataBuilder, PleRecoverSSMMetadataBuilder):
+        assert (
+            builder.get_cudagraph_support(cast(Any, None), cast(Any, None))
+            == AttentionCGSupport.NEVER
+        )
+
+
+@pytest.mark.parametrize("compile_mode", ["VLLM_COMPILE", "NONE"])
+@pytest.mark.parametrize(("breakable", "expected"), [("1", "PIECEWISE"), ("0", "NONE")])
+def test_never_support_falls_back_to_piecewise_with_breakable_cudagraphs(
+    monkeypatch, compile_mode, breakable, expected
+):
+    """A backend without full-graph support must keep piecewise graphs when breakable
+    CUDA graphs split at attention (splitting_ops is empty then, and torch.compile may
+    be off)."""
+    from vllm.config.compilation import (
+        CompilationConfig,
+        CompilationMode,
+        CUDAGraphMode,
+    )
+
+    monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", breakable)
+    compilation_config = CompilationConfig(
+        mode=CompilationMode[compile_mode],
+        cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,
+        splitting_ops=[],
+    )
+    resolved = compilation_config.resolve_cudagraph_mode_and_sizes(
+        AttentionCGSupport.NEVER,
+        "GDNRecoverSSMAttentionBackend",
+        uniform_decode_query_len=4,
+    )
+    assert resolved.name == expected

tests/kernels/mamba/test_recoverssm_gdn.py

@@ -0,0 +1,602 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GDN RecoverSSM: verify against the native fused kernel, commit against the native
+per-token states, and the boundary checks (malformed metadata must raise ValueError
+before any kernel reads or writes out of bounds)."""
+
+from typing import Any, cast
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.mamba.gdn.recoverssm_gdn import (
+    GDNRecoverSSMCommitContext,
+    gdn_recoverssm_verify,
+)
+from vllm.third_party.flash_linear_attention.ops.fused_sigmoid_gating import (
+    fused_sigmoid_gating_delta_rule_update,
+)
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+H, HV, K, V, T = 16, 48, 128, 128, 4
+QLENS = (4, 4, 3)
+
+
+def _inputs(nb: int = 10) -> dict:
+    torch.manual_seed(0)
+    dev = "cuda"
+    tot = sum(QLENS)
+    qsl = torch.tensor([0, 4, 8, 11], dtype=torch.int32, device=dev)
+    return dict(
+        A_log=(torch.rand(HV, device=dev) * 2 - 1).float(),
+        a=torch.randn(tot, HV, device=dev, dtype=torch.bfloat16),
+        b=torch.randn(tot, HV, device=dev, dtype=torch.bfloat16),
+        dt_bias=(torch.randn(HV, device=dev) * 0.5).bfloat16(),
+        q=torch.randn(1, tot, H, K, device=dev, dtype=torch.bfloat16),
+        k=torch.randn(1, tot, H, K, device=dev, dtype=torch.bfloat16),
+        v=torch.randn(1, tot, HV, V, device=dev, dtype=torch.bfloat16),
+        checkpoint_state=(torch.randn(nb, HV, V, K, device=dev) * 0.05).float(),
+        replay_cache=torch.zeros(nb, HV, T, V + K + 1, device=dev),
+        query_start_loc=qsl,
+        state_indices=torch.tensor([2, 5, 7], dtype=torch.int32, device=dev),
+        spec_query_len=T,
+    )
+
+
+def _verify(inp: dict, **override):
+    d = {**inp, **override}
+    return gdn_recoverssm_verify(
+        d["A_log"],
+        d["a"],
+        d["b"],
+        d["dt_bias"],
+        d["q"],
+        d["k"],
+        d["v"],
+        checkpoint_state=d["checkpoint_state"],
+        replay_cache=d["replay_cache"],
+        query_start_loc=d["query_start_loc"],
+        state_indices=d["state_indices"],
+        spec_query_len=d["spec_query_len"],
+        out=d.get("out"),
+    )
+
+
+def _native(inp: dict):
+    return fused_sigmoid_gating_delta_rule_update(
+        A_log=inp["A_log"],
+        a=inp["a"],
+        b=inp["b"],
+        dt_bias=inp["dt_bias"],
+        q=inp["q"],
+        k=inp["k"],
+        v=inp["v"],
+        initial_state=inp["checkpoint_state"].clone(),
+        inplace_final_state=False,
+        cu_seqlens=inp["query_start_loc"].long(),
+        ssm_state_indices=inp["state_indices"].long(),
+        use_qk_l2norm_in_kernel=True,
+    )
+
+
+def _context(inp: dict, ckpt: torch.Tensor) -> GDNRecoverSSMCommitContext:
+    nb = ckpt.shape[0]
+    conv = [torch.zeros(nb, 8, 3 + T - 1, device="cuda", dtype=torch.bfloat16)]
+    return GDNRecoverSSMCommitContext.from_tensors(
+        conv, [ckpt], [inp["replay_cache"]], spec_query_len=T, max_num_reqs=8
+    )
+
+
+def test_verify_matches_native_and_keeps_the_checkpoint():
+    inp = _inputs()
+    before = inp["checkpoint_state"].clone()
+    out = _verify(inp)
+    ref, _ = _native(inp)
+    torch.testing.assert_close(
+        out.float(), ref.reshape(out.shape).float(), rtol=0, atol=0
+    )
+    assert torch.equal(inp["checkpoint_state"], before)
+
+
+@pytest.mark.parametrize("accepted", [1, 2, 3, 4])
+def test_commit_matches_native_per_token_state(accepted):
+    inp = _inputs()
+    _verify(inp)
+    _, states = _native(inp)
+    ckpt = inp["checkpoint_state"].clone()
+    ctx = _context(inp, ckpt)
+    acc = torch.tensor(
+        [min(accepted, q) for q in QLENS], dtype=torch.int32, device="cuda"
+    )
+    ctx.commit(acc, inp["state_indices"], inp["query_start_loc"])
+    for i, blk in enumerate(inp["state_indices"].tolist()):
+        tok = int(inp["query_start_loc"][i]) + int(acc[i]) - 1
+        torch.testing.assert_close(ckpt[blk], states[tok], rtol=1e-5, atol=1e-6)
+    others = [b for b in range(ckpt.shape[0]) if b not in inp["state_indices"].tolist()]
+    assert torch.equal(ckpt[others], inp["checkpoint_state"][others])
+
+
+@pytest.mark.parametrize(
+    "case", ["short_qsl", "capacity", "hv_lt_h", "a_log", "out", "replay", "device"]
+)
+def test_verify_rejects_malformed_inputs(case):
+    inp = _inputs()
+    over, match = {
+        "short_qsl": (
+            dict(query_start_loc=inp["query_start_loc"][:-1]),
+            "query metadata",
+        ),
+        "capacity": (
+            dict(
+                state_indices=inp["state_indices"][:2],
+                query_start_loc=inp["query_start_loc"][:3],
+            ),
+            "activation capacity",
+        ),
+        "hv_lt_h": (
+            dict(
+                v=inp["v"][:, :, :8].contiguous(),
+                a=inp["a"][:, :8].contiguous(),
+                b=inp["b"][:, :8].contiguous(),
+            ),
+            "positive multiple",
+        ),
+        "a_log": (dict(A_log=torch.zeros(HV + 1, device="cuda")), "A_log"),
+        "out": (
+            dict(
+                out=torch.empty(1, 11, HV, V + 1, device="cuda", dtype=torch.bfloat16)
+            ),
+            "output shape",
+        ),
+        "replay": (
+            dict(replay_cache=torch.zeros(10, HV, T, V + K, device="cuda")),
+            "replay buffer",
+        ),
+        "device": (dict(A_log=inp["A_log"].cpu()), "same device"),
+    }[case]
+    with pytest.raises(ValueError, match=match):
+        _verify(inp, **over)
+
+
+def test_commit_rejects_partial_or_short_metadata():
+    inp = _inputs()
+    ctx = _context(inp, inp["checkpoint_state"].clone())
+    acc = torch.ones(3, dtype=torch.int32, device="cuda")
+    bt = torch.zeros(3, 4, dtype=torch.int32, device="cuda")
+    with pytest.raises(ValueError, match="align metadata is incomplete"):
+        ctx.commit(
+            acc,
+            inp["state_indices"],
+            inp["query_start_loc"],
+            block_table=bt,
+            mamba_block_size=16,
+        )
+    with pytest.raises(ValueError, match="commit metadata"):
+        ctx.commit(acc, inp["state_indices"], inp["query_start_loc"][:-1])
+    with pytest.raises(ValueError, match="request mapping"):
+        ctx.commit(
+            acc,
+            inp["state_indices"],
+            inp["query_start_loc"],
+            request_indices=torch.zeros(2, dtype=torch.int32, device="cuda"),
+        )
+
+
+def test_context_rejects_inconsistent_layers():
+    inp = _inputs()
+    ck = inp["checkpoint_state"]
+    conv = torch.zeros(ck.shape[0], 8, 3 + T - 1, device="cuda", dtype=torch.bfloat16)
+    with pytest.raises(ValueError, match="at least one layer"):
+        GDNRecoverSSMCommitContext.from_tensors(
+            [], [], [], spec_query_len=T, max_num_reqs=8
+        )
+    with pytest.raises(ValueError, match="differ"):
+        GDNRecoverSSMCommitContext.from_tensors(
+            [conv], [ck, ck], [inp["replay_cache"]], spec_query_len=T, max_num_reqs=8
+        )
+    with pytest.raises(ValueError, match="block count"):
+        GDNRecoverSSMCommitContext.from_tensors(
+            [conv[:5]], [ck], [inp["replay_cache"]], spec_query_len=T, max_num_reqs=8
+        )
+
+
+# ------------------------------------------------------------------ review round 2:
+# mapping, align, extreme gates
+from vllm.v1.attention.backends.recoverssm_metadata import (  # noqa: E402
+    recoverssm_request_indices,
+    recoverssm_spec_rows,
+)
+
+# (A_log range, a scale, b offset): drive g = -exp(A_log) * softplus(a + dt_bias)
+# and beta = sigmoid(b)
+GATES = {
+    "typical": (0.0, 1.0, 0.0),
+    "g_near_0": (-12.0, 1.0, 0.0),
+    "g_very_negative": (2.5, 4.0, 0.0),
+    "beta_near_0": (0.0, 1.0, -12.0),
+    "beta_near_1": (0.0, 1.0, 12.0),
+}
+
+
+def _make(qlens, nb=16, seed=0, gate="typical"):
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    rnd = lambda *s: torch.randn(*s, device="cuda", generator=g)  # noqa: E731
+    a_log_shift, ascale, boff = GATES[gate]
+    tot = sum(qlens)
+    qsl = torch.tensor(
+        [0, *torch.tensor(qlens).cumsum(0).tolist()], dtype=torch.int32, device="cuda"
+    )
+    return dict(
+        A_log=(
+            torch.rand(HV, device="cuda", generator=g) * 2 - 1 + a_log_shift
+        ).float(),
+        a=(rnd(tot, HV) * ascale).bfloat16(),
+        b=(rnd(tot, HV) + boff).bfloat16(),
+        dt_bias=(rnd(HV) * 0.5).bfloat16(),
+        q=rnd(1, tot, H, K).bfloat16(),
+        k=rnd(1, tot, H, K).bfloat16(),
+        v=rnd(1, tot, HV, V).bfloat16(),
+        checkpoint_state=(rnd(nb, HV, V, K) * 0.05).float(),
+        replay_cache=torch.zeros(nb, HV, T, V + K + 1, device="cuda"),
+        query_start_loc=qsl,
+        spec_query_len=T,
+    )
+
+
+def test_spec_rows_mixed_batch():
+    # rows: 0 spec K=3 | 1 prefill | 2 draft-less decode | 3 spec K=1 | 4 padded
+    # zero-query | 5,6 beyond num_reqs
+    is_prefilling = torch.tensor([False, True, False, False, False, False, False])
+    qsl = torch.tensor([0, 4, 104, 105, 107, 107], dtype=torch.int32)
+    drafts_in = torch.tensor(
+        [3, -1, -1, 1, -1, 2, 0], dtype=torch.int32
+    )  # 5, 6: garbage padding
+    drafts, rows = recoverssm_spec_rows(is_prefilling, qsl, drafts_in, num_reqs=5)
+    assert rows.tolist() == [0, 2, 3]
+    assert drafts.tolist() == [1, -1, 1, 1, -1, -1, -1]
+    idx = recoverssm_request_indices(rows, 3, torch.device("cuda"))
+    assert idx is not None and idx.tolist() == [0, 2, 3]
+    assert (
+        recoverssm_request_indices(torch.tensor([0, 1, 2]), 3, torch.device("cuda"))
+        is None
+    )
+    with pytest.raises(ValueError, match="spec rows classified"):
+        recoverssm_request_indices(rows, 2, torch.device("cuda"))
+
+
+def test_commit_maps_noncontiguous_requests_in_align_mode():
+    """Spec rows 0, 2, 3 of a 5-request batch; accepted counts, block table and
+    num_computed are indexed by request row. Every spec request must commit its own
+    state into its own block."""
+    qlens = [4, 1, 2]  # the spec rows' windows (K=3, draft-less, K=1)
+    inp = _make(qlens, nb=40)
+    bs = 8
+    bt = torch.arange(1, 1 + 5 * 6, dtype=torch.int32, device="cuda").reshape(
+        5, 6
+    )  # request row -> blocks
+    nc = torch.tensor(
+        [5, 0, 16, 9, 0], dtype=torch.int32, device="cuda"
+    )  # per request row
+    req_rows = torch.tensor([0, 2, 3], dtype=torch.int32, device="cuda")
+    src = torch.stack([bt[r, int(nc[r]) // bs] for r in req_rows.tolist()]).contiguous()
+    inp["state_indices"] = src
+    _verify(inp)
+    _, states = _native(inp)
+    ckpt = inp["checkpoint_state"].clone()
+    ctx = _context(inp, ckpt)
+    acc = torch.tensor(
+        [3, 9, 1, 2, 9], dtype=torch.int32, device="cuda"
+    )  # rows 1, 4 must be ignored
+    ctx.commit(
+        acc,
+        src,
+        inp["query_start_loc"],
+        request_indices=req_rows,
+        block_table=bt,
+        num_computed_tokens=nc,
+        mamba_block_size=bs,
+    )
+    for i, r in enumerate(req_rows.tolist()):
+        n = min(int(acc[r]), qlens[i])
+        c0 = int(nc[r])
+        final_blk = int(bt[r, max(c0 + n - 1, 0) // bs])
+        tok = int(inp["query_start_loc"][i]) + n - 1
+        torch.testing.assert_close(ckpt[final_blk], states[tok], rtol=1e-5, atol=1e-6)
+
+
+# (num_computed, accepted, block size): boundary after 1 / after 2 / ending exactly on
+# it / crossing and continuing / accepted = 1 without crossing / the full window
+# crossing after 3. A window ending exactly on a boundary keeps its state in the block
+# holding its last token: the next block may still be unallocated.
+ALIGN_CASES = [
+    (3, 2, 4),
+    (2, 3, 4),
+    (0, 4, 4),
+    (6, 4, 4),
+    (1, 1, 4),
+    (5, 4, 8),
+    (7, 4, 8),
+    (4, 4, 8),
+]
+
+
+@pytest.mark.parametrize("nc,accepted,bs", ALIGN_CASES)
+def test_align_boundary_and_final_state_match_native(nc, accepted, bs):
+    inp = _make([T], nb=12)
+    bt = torch.arange(2, 8, dtype=torch.int32, device="cuda").reshape(1, 6)
+    src = bt[0, nc // bs].reshape(1).contiguous()
+    inp["state_indices"] = src
+    _verify(inp)
+    _, states = _native(inp)
+    ckpt = inp["checkpoint_state"].clone()
+    ctx = _context(inp, ckpt)
+    ctx.commit(
+        torch.tensor([accepted], dtype=torch.int32, device="cuda"),
+        src,
+        inp["query_start_loc"],
+        block_table=bt,
+        num_computed_tokens=torch.tensor([nc], dtype=torch.int32, device="cuda"),
+        mamba_block_size=bs,
+    )
+    next_boundary = (nc // bs + 1) * bs
+    final_blk = int(bt[0, max(nc + accepted - 1, 0) // bs])
+    if (
+        nc + accepted >= next_boundary
+    ):  # crosses: the boundary state goes into the source block
+        torch.testing.assert_close(
+            ckpt[int(src)], states[next_boundary - nc - 1], rtol=1e-5, atol=1e-6
+        )
+    torch.testing.assert_close(
+        ckpt[final_blk], states[accepted - 1], rtol=1e-5, atol=1e-6
+    )
+
+
+@pytest.mark.parametrize("gate", list(GATES))
+@pytest.mark.parametrize("seed", [0, 1, 2])
+def test_commit_matches_native_under_extreme_gates(gate, seed):
+    inp = _make(list(QLENS), nb=10, seed=seed, gate=gate)
+    inp["state_indices"] = torch.tensor([2, 5, 7], dtype=torch.int32, device="cuda")
+    _verify(inp)
+    _, states = _native(inp)
+    for accepted in (1, 2, 3, 4):
+        ckpt = inp["checkpoint_state"].clone()
+        ctx = _context(inp, ckpt)
+        acc = torch.tensor(
+            [min(accepted, q) for q in QLENS], dtype=torch.int32, device="cuda"
+        )
+        ctx.commit(acc, inp["state_indices"], inp["query_start_loc"])
+        for i, blk in enumerate(inp["state_indices"].tolist()):
+            tok = int(inp["query_start_loc"][i]) + int(acc[i]) - 1
+            torch.testing.assert_close(ckpt[blk], states[tok], rtol=1e-5, atol=1e-6)
+
+
+@pytest.mark.parametrize("batch", [1, 2, 4])
+def test_verify_on_strided_views_writes_in_place(batch):
+    """The layer passes views of the packed conv output (token stride = the packed
+    width) and the layer's output buffer; result and replay record must equal the
+    contiguous-copy path."""
+    torch.manual_seed(batch)
+    H, HV, K, V, SQ, NB = 16, 48, 128, 128, 4, 6
+    T = batch * SQ
+    qd, vd = H * K, HV * V
+    packed = torch.randn(T, 2 * qd + 2 * vd, device="cuda", dtype=torch.bfloat16)
+    mixed = packed[:, : 2 * qd + vd]
+    a = torch.randn(T, HV, device="cuda", dtype=torch.bfloat16)
+    b = torch.randn(T, HV, device="cuda", dtype=torch.bfloat16)
+    A_log = torch.randn(HV, device="cuda")
+    dt_bias = torch.randn(HV, device="cuda")
+    checkpoint = torch.randn(NB, HV, V, K, device="cuda") * 0.05
+    qsl = torch.arange(0, T + 1, SQ, device="cuda", dtype=torch.int32)
+    si = torch.arange(1, batch + 1, device="cuda", dtype=torch.int32)
+
+    def run(q, k, v, out=None):
+        replay = torch.zeros(NB, HV, SQ, V + K + 1, device="cuda")
+        o = gdn_recoverssm_verify(
+            A_log,
+            a,
+            b,
+            dt_bias,
+            q,
+            k,
+            v,
+            checkpoint_state=checkpoint,
+            replay_cache=replay,
+            query_start_loc=qsl,
+            state_indices=si,
+            spec_query_len=SQ,
+            out=out,
+        )
+        return o, replay
+
+    q, k, v = (
+        x.contiguous().view(1, T, -1, d)
+        for x, d in zip(torch.split(mixed, [qd, qd, vd], dim=-1), (K, K, V))
+    )
+    want_out, want_replay = run(q, k, v)
+
+    core = torch.zeros(T + 3, HV, V, device="cuda", dtype=torch.bfloat16)
+    views = (
+        mixed[:, :qd].view(1, T, -1, K),
+        mixed[:, qd : 2 * qd].view(1, T, -1, K),
+        mixed[:, 2 * qd :].view(1, T, -1, V),
+    )
+    got_out, got_replay = run(*views, out=core[:T].unsqueeze(0))
+    assert got_out.data_ptr() == core.data_ptr()
+    assert torch.equal(core[:T].unsqueeze(0), want_out)
+    assert torch.equal(got_replay, want_replay)
+
+
+def test_verify_on_strided_views_with_variable_query_lengths():
+    """Token-strided views with ragged windows (query lengths 4, 1, 3), as MTP
+    scheduling produces them."""
+    torch.manual_seed(7)
+    H, HV, K, V, SQ, NB = 16, 48, 128, 128, 4, 6
+    qsl_list = [0, 4, 5, 8]
+    T = qsl_list[-1]
+    qd, vd = H * K, HV * V
+    packed = torch.randn(T, 2 * qd + 2 * vd, device="cuda", dtype=torch.bfloat16)
+    mixed = packed[:, : 2 * qd + vd]
+    a = torch.randn(T, HV, device="cuda", dtype=torch.bfloat16)
+    b = torch.randn(T, HV, device="cuda", dtype=torch.bfloat16)
+    A_log = torch.randn(HV, device="cuda")
+    dt_bias = torch.randn(HV, device="cuda")
+    checkpoint = torch.randn(NB, HV, V, K, device="cuda") * 0.05
+    qsl = torch.tensor(qsl_list, device="cuda", dtype=torch.int32)
+    si = torch.tensor([1, 3, 5], device="cuda", dtype=torch.int32)
+
+    def run(q, k, v, out=None):
+        replay = torch.zeros(NB, HV, SQ, V + K + 1, device="cuda")
+        o = gdn_recoverssm_verify(
+            A_log,
+            a,
+            b,
+            dt_bias,
+            q,
+            k,
+            v,
+            checkpoint_state=checkpoint,
+            replay_cache=replay,
+            query_start_loc=qsl,
+            state_indices=si,
+            spec_query_len=SQ,
+            out=out,
+        )
+        return o, replay
+
+    q, k, v = (
+        x.contiguous().view(1, T, -1, d)
+        for x, d in zip(torch.split(mixed, [qd, qd, vd], dim=-1), (K, K, V))
+    )
+    want_out, want_replay = run(q, k, v)
+    core = torch.zeros(T + 2, HV, V, device="cuda", dtype=torch.bfloat16)
+    views = (
+        mixed[:, :qd].view(1, T, -1, K),
+        mixed[:, qd : 2 * qd].view(1, T, -1, K),
+        mixed[:, 2 * qd :].view(1, T, -1, V),
+    )
+    _, got_replay = run(*views, out=core[:T].unsqueeze(0))
+    assert torch.equal(core[:T].unsqueeze(0), want_out)
+    assert torch.equal(got_replay, want_replay)
+
+
+def test_layer_verify_qkv_views_and_fallback():
+    """The layer's guard: views for a [tokens, q|k|v] conv output with a contiguous
+    last dimension, the contiguous-copy fallback for any other layout; same values
+    either way."""
+    from types import SimpleNamespace
+
+    from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
+        QwenGatedDeltaNetAttention as Layer,
+    )
+
+    H, HV, K, V, T = 4, 8, 16, 16, 5
+    layer = SimpleNamespace(
+        key_dim=H * K, value_dim=HV * V, tp_size=1, head_k_dim=K, head_v_dim=V
+    )
+    layer.rearrange_mixed_qkv = lambda x: Layer.rearrange_mixed_qkv(cast(Any, layer), x)
+    width = 2 * H * K + HV * V
+    packed = torch.randn(T, width + HV * V, dtype=torch.bfloat16)  # q|k|v|z, like qkvz
+    strided = packed[:, :width]
+    q, k, v, is_view = Layer._recoverssm_verify_qkv(cast(Any, layer), strided)
+    assert is_view
+    assert q.data_ptr() == packed.data_ptr()
+    rq, rk, rv = Layer.rearrange_mixed_qkv(cast(Any, layer), strided)
+    assert torch.equal(q, rq) and torch.equal(k, rk) and torch.equal(v, rv)
+
+    # Column-major storage: last dimension not contiguous -> copies, same values.
+    col_major = strided.t().contiguous().t()
+    assert col_major.stride(-1) != 1
+    q2, k2, v2, is_view2 = Layer._recoverssm_verify_qkv(cast(Any, layer), col_major)
+    assert not is_view2
+    assert torch.equal(q2, rq) and torch.equal(k2, rk) and torch.equal(v2, rv)
+
+
+# ------------------------------------------------------------- BF16 recurrent state
+# RecoverSSM upcasts the BF16 checkpoint, carries the accepted window in FP32 and rounds
+# once when it writes the committed state. The reference is the FP32 trajectory from the
+# same (already BF16-rounded) checkpoint; the committed state must equal its BF16
+# rounding up to one BF16 ulp, and be no further from it than the native path is.
+BF16_ULP = 2**-7
+
+
+def _bf16_case(qlens, nb=12, seed=0):
+    inp = _make(qlens, nb=nb, seed=seed)
+    inp["checkpoint_state"] = inp["checkpoint_state"].bfloat16()
+    ref = dict(inp)
+    ref["checkpoint_state"] = inp["checkpoint_state"].float()
+    return inp, ref
+
+
+def test_verify_bf16_state_matches_native():
+    inp, _ = _bf16_case(list(QLENS))
+    inp["state_indices"] = torch.tensor([2, 5, 7], dtype=torch.int32, device="cuda")
+    before = inp["checkpoint_state"].clone()
+    out = _verify(inp)
+    ref, _ = _native(inp)
+    torch.testing.assert_close(
+        out.float(), ref.reshape(out.shape).float(), rtol=0, atol=0
+    )
+    assert torch.equal(inp["checkpoint_state"], before)
+
+
+@pytest.mark.parametrize("accepted", [1, 2, 3, 4])
+@pytest.mark.parametrize("seed", [0, 1])
+def test_commit_bf16_state_rounds_once(accepted, seed):
+    inp, ref_inp = _bf16_case(list(QLENS), seed=seed)
+    idx = torch.tensor([2, 5, 7], dtype=torch.int32, device="cuda")
+    inp["state_indices"] = ref_inp["state_indices"] = idx
+    _verify(inp)
+    _, native_states = _native(inp)
+    _, ref_states = _native(ref_inp)
+    ckpt = inp["checkpoint_state"].clone()
+    _context(inp, ckpt).commit(
+        torch.tensor(
+            [min(accepted, q) for q in QLENS], dtype=torch.int32, device="cuda"
+        ),
+        idx,
+        inp["query_start_loc"],
+    )
+    for i, blk in enumerate(idx.tolist()):
+        tok = int(inp["query_start_loc"][i]) + min(accepted, QLENS[i]) - 1
+        ref = ref_states[tok].float()
+        got = ckpt[blk].float()
+        torch.testing.assert_close(
+            got, ref.bfloat16().float(), rtol=BF16_ULP, atol=1e-5
+        )
+        err_rssm = (got - ref).abs().mean()
+        err_native = (native_states[tok].bfloat16().float() - ref).abs().mean()
+        assert err_rssm <= err_native * 1.05 + 1e-7, (err_rssm, err_native)
+
+
+@pytest.mark.parametrize("nc,accepted,bs", ALIGN_CASES)
+def test_align_boundary_bf16_state(nc, accepted, bs):
+    inp, ref_inp = _bf16_case([T])
+    bt = torch.arange(2, 8, dtype=torch.int32, device="cuda").reshape(1, 6)
+    src = bt[0, nc // bs].reshape(1).contiguous()
+    inp["state_indices"] = ref_inp["state_indices"] = src
+    _verify(inp)
+    _, ref_states = _native(ref_inp)
+    ckpt = inp["checkpoint_state"].clone()
+    _context(inp, ckpt).commit(
+        torch.tensor([accepted], dtype=torch.int32, device="cuda"),
+        src,
+        inp["query_start_loc"],
+        block_table=bt,
+        num_computed_tokens=torch.tensor([nc], dtype=torch.int32, device="cuda"),
+        mamba_block_size=bs,
+    )
+    next_boundary = (nc // bs + 1) * bs
+    final_blk = int(bt[0, max(nc + accepted - 1, 0) // bs])
+    if nc + accepted >= next_boundary:
+        torch.testing.assert_close(
+            ckpt[int(src)].float(),
+            ref_states[next_boundary - nc - 1].float().bfloat16().float(),
+            rtol=BF16_ULP,
+            atol=1e-5,
+        )
+    torch.testing.assert_close(
+        ckpt[final_blk].float(),
+        ref_states[accepted - 1].float().bfloat16().float(),
+        rtol=BF16_ULP,
+        atol=1e-5,
+    )

tests/kernels/mamba/test_recoverssm_ple.py

@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""PLE short-conv RecoverSSM against the sequential reference.
+
+Reference: the stock decode path fed the accepted tokens one at a time.
+RecoverSSM: one spec-mode conv over the whole window from the checkpoint
+(offset 0), then ``_PleConvCommit`` compacts the accepted suffix to the front.
+The conv history is pure data movement, so the states must be bit-identical,
+and the outputs of the accepted tokens must match as well.
+"""
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
+from vllm.models.qwen4_exp.nvidia.ops.ple import ple_conv
+from vllm.v1.attention.backends.ple_recoverssm import _PleConvCommit
+
+C, KSIZE, DIL, SQ, NB = 512, 4, 3, 4, 8
+STATE_LEN = (KSIZE - 1) * DIL
+WIDTH = STATE_LEN + SQ - 1
+DEV = "cuda"
+
+
+def _state():
+    if is_conv_state_dim_first():
+        raw = torch.zeros(NB, C, WIDTH, dtype=torch.bfloat16, device=DEV)
+        return raw, raw
+    raw = torch.zeros(NB, WIDTH, C, dtype=torch.bfloat16, device=DEV)
+    return raw, raw.transpose(-1, -2)
+
+
+def _i32(xs):
+    return torch.tensor(xs, dtype=torch.int32, device=DEV)
+
+
+def _setup(seed, blocks):
+    g = torch.Generator(device=DEV).manual_seed(seed)
+    w = torch.randn(C, KSIZE, generator=g, device=DEV).to(torch.bfloat16)
+    x = torch.randn(len(blocks) * SQ, C, generator=g, device=DEV).to(torch.bfloat16)
+    outer = torch.randn_like(x)
+    raw, view = _state()
+    for b in blocks:
+        view[b, :, :STATE_LEN] = torch.randn(C, STATE_LEN, generator=g, device=DEV).to(
+            torch.bfloat16
+        )
+    return w, x, outer, raw, view
+
+
+def _sequential(w, x, outer, view, block, first, count):
+    """Stock decode, one token at a time, on a copy of the checkpoint block."""
+    ref = view.clone()
+    res = torch.zeros(count, C, dtype=torch.bfloat16, device=DEV)
+    for i in range(count):
+        t = first + i
+        ple_conv(
+            inputs=x[t : t + 1],
+            residual=res[i : i + 1],
+            conv_state=ref,
+            conv_weights=w,
+            state_indices=_i32([block]),
+            outer_residual=outer[t : t + 1],
+            mode="decode",
+            dilation=DIL,
+        )
+    return ref[block, :, :STATE_LEN].clone(), res
+
+
+def _spec(w, x, outer, view, blocks):
+    res = torch.zeros(len(blocks) * SQ, C, dtype=torch.bfloat16, device=DEV)
+    ple_conv(
+        inputs=x,
+        residual=res,
+        conv_state=view,
+        conv_weights=w,
+        state_indices=_i32(blocks),
+        outer_residual=outer,
+        mode="spec",
+        dilation=DIL,
+        query_start_loc=_i32(list(range(0, len(blocks) * SQ + 1, SQ))),
+        num_accepted_tokens=_i32([1] * len(blocks)),
+        spec_query_len=SQ,
+    )
+    return res
+
+
+@pytest.mark.parametrize("accepted", [1, 2, SQ])
+@pytest.mark.parametrize("seed", [0, 1])
+def test_ple_commit_matches_sequential(accepted, seed):
+    block = 3
+    w, x, outer, raw, view = _setup(seed, [block])
+    want_state, want_out = _sequential(w, x, outer, view, block, 0, accepted)
+    history = torch.cat([view[block, :, :STATE_LEN], x[:accepted].T], dim=1)
+
+    got_out = _spec(w, x, outer, view, [block])
+    ctx = _PleConvCommit([raw], spec_query_len=SQ, max_num_reqs=NB)
+    ctx.commit(_i32([accepted]), _i32([block]), _i32([0, SQ]))
+
+    got_state = view[block, :, :STATE_LEN]
+    assert torch.equal(got_state, want_state)
+    assert torch.equal(got_state, history[:, -STATE_LEN:])
+    assert torch.equal(got_out[:accepted], want_out)
+
+
+@pytest.mark.parametrize(
+    ("num_computed", "accepted", "crosses"),
+    [(6, 3, True), (6, 2, True), (1, 2, False), (7, SQ, True)],
+)
+def test_ple_commit_align_boundary(num_computed, accepted, crosses):
+    block_size, src, nxt = 8, 2, 5
+    w, x, outer, raw, view = _setup(7, [src])
+    boundary_len = block_size - num_computed
+    want_final, _ = _sequential(w, x, outer, view, src, 0, accepted)
+    want_boundary, _ = _sequential(w, x, outer, view, src, 0, min(boundary_len, SQ))
+
+    _spec(w, x, outer, view, [src])
+    ctx = _PleConvCommit([raw], spec_query_len=SQ, max_num_reqs=NB)
+    ctx.commit(
+        _i32([accepted]),
+        _i32([src]),
+        _i32([0, SQ]),
+        block_table=_i32([[src, nxt]]),
+        num_computed_tokens=_i32([num_computed]),
+        mamba_block_size=block_size,
+    )
+
+    final_block = nxt if num_computed + accepted > block_size else src
+    assert torch.equal(view[final_block, :, :STATE_LEN], want_final)
+    if crosses:
+        assert torch.equal(view[src, :, :STATE_LEN], want_boundary)
+
+
+def test_ple_commit_non_contiguous_request_rows():
+    blocks, accepted = [1, 6], [3, 1]
+    w, x, outer, raw, view = _setup(11, blocks)
+    want = [
+        _sequential(w, x, outer, view, b, i * SQ, a)[0]
+        for i, (b, a) in enumerate(zip(blocks, accepted))
+    ]
+
+    _spec(w, x, outer, view, blocks)
+    # Spec rows 0 and 1 belong to requests 1 and 3 of a four-request batch.
+    counts = _i32([9, accepted[0], 9, accepted[1]])
+    ctx = _PleConvCommit([raw], spec_query_len=SQ, max_num_reqs=NB)
+    ctx.commit(
+        counts,
+        _i32(blocks),
+        _i32([0, SQ, 2 * SQ]),
+        request_indices=_i32([1, 3]),
+    )
+
+    for b, ref in zip(blocks, want):
+        assert torch.equal(view[b, :, :STATE_LEN], ref)

tests/model_executor/offloader/test_prefetch.py

@@ -2,8 +2,13 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import pytest
+import torch
+import torch.nn as nn
 
-from vllm.model_executor.offloader.prefetch import _get_next_prefetch_index
+from vllm.model_executor.offloader.prefetch import (
+    _get_next_prefetch_index,
+    _TowerUnitStreamer,
+)
 
 pytestmark = pytest.mark.skip_global_cleanup
 
@@ -35,3 +40,61 @@ def test_next_prefetch_index_preserves_slot_ownership(
             target % prefetch_step == index % prefetch_step
             for index, target in enumerate(targets)
         )
+
+
+def test_tower_unit_streamer_roundtrip_cpu(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Host-resident weights must compute identically to resident weights,
+    and be back on host storage after each forward."""
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+
+    lin = nn.Linear(4, 4, bias=False)
+    weight = lin.weight.detach().clone()
+
+    streamer = _TowerUnitStreamer(lin, ["weight"], torch.device("cpu"))
+    streamer.offload_to_host_now()
+    assert lin.weight.device.type == "cpu"
+
+    streamer.sync_storage()
+    assert streamer.offloaded_bytes == weight.numel() * weight.element_size()
+
+    streamer.hook_forward()
+    x = torch.randn(2, 4)
+    out = lin(x)
+    assert torch.allclose(out, x @ weight.T)
+    # Weights released back to host after the forward.
+    assert lin.weight.device.type == "cpu"
+
+    # Re-running post_init-style hooking must not double-wrap the forward.
+    streamer.hook_forward()
+    assert torch.allclose(lin(x), out)
+
+
+def test_offload_config_tower_only_prefetch_is_valid() -> None:
+    from vllm.config.offload import OffloadConfig, PrefetchOffloadConfig
+
+    cfg = OffloadConfig(
+        offload_backend="prefetch",
+        prefetch=PrefetchOffloadConfig(offload_towers=True),
+    )
+    assert cfg.prefetch.offload_towers
+    assert cfg.prefetch.offload_group_size == 0
+
+
+def test_offload_config_prefetch_requires_layers_or_towers() -> None:
+    from vllm.config.offload import OffloadConfig
+
+    with pytest.raises(ValueError):
+        OffloadConfig(offload_backend="prefetch")
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_create_offloader_selects_prefetch_for_towers() -> None:
+    from vllm.config.offload import OffloadConfig, PrefetchOffloadConfig
+    from vllm.model_executor.offloader.base import create_offloader
+    from vllm.model_executor.offloader.prefetch import PrefetchOffloader
+
+    offloader = create_offloader(
+        OffloadConfig(prefetch=PrefetchOffloadConfig(offload_towers=True))
+    )
+    assert isinstance(offloader, PrefetchOffloader)
+    assert offloader.supports_tower_offload

tests/model_executor/test_qwen_triton_warmup.py

@@ -1,9 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
+import vllm.model_executor.warmup.qwen_triton_warmup as qwen_triton_warmup_module
 from vllm.model_executor.warmup.qwen_triton_warmup import (
     _FLA_POST_CONV_WARMUP_LENGTHS,
     _QwenGDNWarmupConfig,
@@ -53,3 +56,57 @@ def test_qwen_gdn_prefill_warmup_kernels_compile_on_gpu() -> None:
     _warm_fused_post_conv_kernel(device, config)
     assert _FLA_POST_CONV_WARMUP_LENGTHS == (1, 2, 16)
     torch.accelerator.synchronize(device)
+
+
+def _warmup_runner() -> SimpleNamespace:
+    return SimpleNamespace(
+        device=torch.device("cpu"),
+        compilation_config=SimpleNamespace(static_forward_context={}),
+        max_num_tokens=1,
+        is_pooling_model=False,
+    )
+
+
+def _warmup_model_config(model_type: str) -> SimpleNamespace:
+    return SimpleNamespace(
+        hf_text_config=SimpleNamespace(model_type=model_type),
+        hf_config=SimpleNamespace(model_type=model_type),
+        dtype=torch.bfloat16,
+    )
+
+
+@pytest.mark.parametrize(
+    ("model_type", "dispatched"),
+    [
+        ("qwen4_exp", True),
+        ("qwen4_exp_text", True),
+        ("qwen3_next", True),
+        # The drafter never reaches this warmup: the worker passes its own
+        # model_config, which is the target's.
+        ("qwen4_exp_mtp", False),
+        ("llama", False),
+    ],
+)
+def test_qwen_triton_warmup_dispatch_gate(
+    monkeypatch: pytest.MonkeyPatch,
+    model_type: str,
+    dispatched: bool,
+) -> None:
+    """Which model types reach the Qwen Triton warmup at all.
+
+    Stops at the gate: the kernels themselves need a GPU and are covered by the
+    CUDA test above.
+    """
+    reached: list[object] = []
+
+    def _record(static_forward_context: object) -> None:
+        reached.append(static_forward_context)
+        return None
+
+    monkeypatch.setattr(qwen_triton_warmup_module, "_qwen_gdn_warmup_config", _record)
+
+    qwen_triton_warmup_module.qwen_triton_warmup(
+        _warmup_runner(), _warmup_model_config(model_type)
+    )
+
+    assert bool(reached) is dispatched

tests/models/glm5next/test_kda_state_dtype.py

@@ -24,7 +24,7 @@ def test_kda_recurrent_state_dtype_follows_ssm_cache_dtype(ssm_dtype, expected):
     cache_config = SimpleNamespace(
         mamba_cache_dtype="auto",
         mamba_ssm_cache_dtype=ssm_dtype,
-        use_kda_recoverssm=False,
+        use_recoverssm=False,
     )
     vllm_config = SimpleNamespace(model_config=model_config, cache_config=cache_config)
     want = (torch.bfloat16, expected)

tests/models/kimi_k3/test_kda.py

@@ -128,7 +128,7 @@ def test_kda_recoverssm_config_state_layout():
         cache_config=SimpleNamespace(
             mamba_cache_dtype="auto",
             mamba_ssm_cache_dtype="auto",
-            use_kda_recoverssm=True,
+            use_recoverssm=True,
         ),
         parallel_config=SimpleNamespace(tensor_parallel_size=1),
         speculative_config=SimpleNamespace(num_speculative_tokens=2),

tests/models/kimi_k3/test_kda_metadata.py

@@ -122,7 +122,7 @@ def _make_builder(
     )
     vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode
     vllm_config.cache_config.use_replayssm = use_recoverssm
-    vllm_config.cache_config.use_kda_recoverssm = use_recoverssm
+    vllm_config.cache_config.use_recoverssm = use_recoverssm
     vllm_config.cache_config.prefix_match_unit = prefix_match_unit
     hash_block_size = hash_block_size or prefix_match_unit or mamba_block_size
     vllm_config.cache_config.hash_block_size = hash_block_size
@@ -169,7 +169,7 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch):
         cache_config=SimpleNamespace(
             mamba_cache_dtype="auto",
             mamba_ssm_cache_dtype="auto",
-            use_kda_recoverssm=True,
+            use_recoverssm=True,
         ),
         parallel_config=SimpleNamespace(tensor_parallel_size=1),
         speculative_config=SimpleNamespace(num_speculative_tokens=2),
@@ -203,7 +203,7 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch):
         ),
         cache_config=SimpleNamespace(
             mamba_cache_mode="align",
-            use_kda_recoverssm=True,
+            use_recoverssm=True,
             prefix_match_unit=None,
         ),
         parallel_config=SimpleNamespace(decode_context_parallel_size=1),

tests/models/kimi_k3/test_recoverssm_boundary.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Align-mode RecoverSSM state columns at recurrent block boundaries.
+
+A speculative window whose accepted tokens end exactly on a recurrent block
+boundary must keep its state in the block holding the last committed token.
+The following block can still be unallocated, so selecting it would drop the
+commit while the bookkeeping recorded the state as advanced.
+"""
+
+import pytest
+import torch
+
+from vllm.models.kimi_k3.nvidia.ops.recoverssm import _prepare_commit_plan_kernel
+from vllm.platforms import current_platform
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
+from vllm.v1.worker.gpu.model_states.recoverssm import (
+    _postprocess_recoverssm_align_kernel,
+)
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cuda_alike(),
+    reason="The RecoverSSM kernels require a CUDA-alike device.",
+)
+
+BLOCK_SIZE = 2048
+SPEC_QUERY_LEN = 4
+BLOCKS = [100, 101, 102, 103, 104, 105]
+
+# (computed before the step, committed tokens, allocated block columns,
+#  final column, boundary column or None)
+CASES = [
+    # 8,188 + 4 ends exactly on the 8,192 boundary; column 4 is unallocated.
+    (8188, 4, 4, 3, 3),
+    # The window crosses the boundary: state moves on, checkpoint stays behind.
+    (8190, 4, 5, 4, 3),
+    # The window stays inside column 3.
+    (8180, 4, 4, 3, None),
+    # The first token after a boundary belongs to the new column.
+    (8192, 1, 5, 4, None),
+    # Single-token commit ending on the first boundary.
+    (2047, 1, 1, 0, 0),
+]
+
+
+@pytest.mark.parametrize(
+    "num_computed,commit_len,allocated,final_col,boundary_col", CASES
+)
+def test_commit_plan_keeps_boundary_state_in_last_computed_block(
+    num_computed, commit_len, allocated, final_col, boundary_col
+):
+    device = current_platform.device_type
+    columns = [
+        block if col < allocated else NULL_BLOCK_ID for col, block in enumerate(BLOCKS)
+    ]
+    block_table = torch.tensor([columns], dtype=torch.int32, device=device)
+    num_accepted = torch.tensor([commit_len], dtype=torch.int32, device=device)
+    state_indices = torch.tensor(
+        [BLOCKS[(num_computed - 1) // BLOCK_SIZE]], dtype=torch.int32, device=device
+    )
+    query_start_loc = torch.tensor(
+        [0, SPEC_QUERY_LEN], dtype=torch.int32, device=device
+    )
+    computed = torch.tensor([num_computed], dtype=torch.int32, device=device)
+    commit_lens, final_indices, boundary_indices, recovery_lens = (
+        torch.full((1,), -7, dtype=torch.int32, device=device) for _ in range(4)
+    )
+
+    _prepare_commit_plan_kernel[(1,)](
+        num_accepted,
+        None,
+        state_indices,
+        query_start_loc,
+        block_table,
+        computed,
+        commit_lens,
+        final_indices,
+        boundary_indices,
+        recovery_lens,
+        NULL_BLOCK_ID,
+        BLOCK_SIZE,
+        block_table.shape[1],
+        num_accepted.stride(0),
+        0,
+        state_indices.stride(0),
+        query_start_loc.stride(0),
+        block_table.stride(0),
+        block_table.stride(1),
+        computed.stride(0),
+        SPEC_QUERY_LEN=SPEC_QUERY_LEN,
+        num_warps=1,
+    )
+
+    assert commit_lens.item() == commit_len
+    assert final_indices.item() == BLOCKS[final_col]
+    expected_boundary = NULL_BLOCK_ID if boundary_col is None else BLOCKS[boundary_col]
+    assert boundary_indices.item() == expected_boundary
+    if boundary_col is not None:
+        next_boundary = (num_computed // BLOCK_SIZE + 1) * BLOCK_SIZE
+        assert recovery_lens.item() == next_boundary - num_computed
+
+
+@pytest.mark.parametrize(
+    "num_computed,commit_len,allocated,final_col,boundary_col", CASES
+)
+def test_postprocess_records_the_committed_state_column(
+    num_computed, commit_len, allocated, final_col, boundary_col
+):
+    device = current_platform.device_type
+    idx_mapping = torch.tensor([2], dtype=torch.int32, device=device)
+    num_sampled = torch.tensor([commit_len], dtype=torch.int32, device=device)
+    computed = torch.tensor([num_computed], dtype=torch.int32, device=device)
+    state_idx = torch.full((4,), -7, dtype=torch.int32, device=device)
+    num_accepted = torch.full((4,), -7, dtype=torch.int32, device=device)
+
+    _postprocess_recoverssm_align_kernel[(1,)](
+        idx_mapping,
+        num_sampled,
+        None,
+        computed,
+        state_idx,
+        num_accepted,
+        MAMBA_BLOCK_SIZE=BLOCK_SIZE,
+        BLOCK_TABLE_WIDTH=len(BLOCKS),
+    )
+
+    assert state_idx[2].item() == final_col
+    assert num_accepted[2].item() == 1
+    assert state_idx[[0, 1, 3]].tolist() == [-7, -7, -7]

tests/models/qwen4_exp/test_config.py

@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import logging
+from collections.abc import Iterator
+from contextlib import contextmanager
 from importlib import import_module
 from types import SimpleNamespace
 from unittest.mock import patch
@@ -15,10 +18,12 @@ from vllm.model_executor.models.config import (
     Qwen3_5ForConditionalGenerationConfig,
     Qwen4ExpForConditionalGenerationConfig,
 )
+from vllm.models.qwen4_exp.common.mtp import make_mtp_hidden_buffer
 from vllm.models.qwen4_exp.nvidia.model_state import Qwen4ExpModelState
 from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
 
 from ...utils import spawn_new_process_for_each_test
+from .test_ple import _make_file_gather_embedding
 
 
 def _text_config(**kwargs) -> Qwen4ExpTextConfig:
@@ -216,6 +221,7 @@ def test_qwen4_exp_model_state_prepares_ngram_context() -> None:
     model_state.ngram_context = torch.empty((8, 3), dtype=torch.int32)
     model_state.ngram_context_offsets = torch.arange(-3, 0, dtype=torch.int64)
     model_state.ple_query_start_loc = torch.empty(9, dtype=torch.int32)
+    model_state.file_gather_modules = []
 
     input_batch = SimpleNamespace(
         num_reqs=2,
@@ -265,6 +271,39 @@ def test_qwen4_exp_model_state_prepares_ngram_context() -> None:
     assert model_inputs["ngram_context"].data_ptr() == ngram_context.data_ptr()
 
 
+def test_qwen4_exp_model_state_stages_file_gather_rows(monkeypatch, tmp_path) -> None:
+    """Rows are staged from the same PLE inputs the forward receives."""
+    module, table = _make_file_gather_embedding(monkeypatch, tmp_path, bind=True)
+    model_state = object.__new__(Qwen4ExpModelState)
+    model_state.uses_ngram_embedding = True
+    model_state.ngram_eos_token_id = 0
+    model_state.ngram_context = torch.empty((4, 1), dtype=torch.int32)
+    model_state.ngram_context_offsets = torch.arange(-1, 0, dtype=torch.int64)
+    model_state.ple_query_start_loc = torch.empty(5, dtype=torch.int32)
+    model_state.file_gather_modules = [module]
+    input_batch = SimpleNamespace(
+        num_reqs=2,
+        num_reqs_after_padding=2,
+        idx_mapping=torch.tensor([0, 1]),
+        query_start_loc=torch.tensor([0, 4, 6], dtype=torch.int32),
+        input_ids=torch.tensor([5, 9, 3, 7, 2, 11]),
+    )
+    req_states = SimpleNamespace(
+        num_computed_tokens=SimpleNamespace(gpu=torch.tensor([2, 1])),
+        all_token_ids=SimpleNamespace(gpu=torch.tensor([[1, 8, 6], [4, 3, 3]])),
+    )
+
+    with patch.object(MambaHybridModelState, "prepare_inputs", return_value={}):
+        model_inputs = model_state.prepare_inputs(input_batch, req_states)
+
+    ids = module.compute_ngram_ids(
+        input_batch.input_ids,
+        model_inputs["query_start_loc"],
+        model_inputs["ngram_context"],
+    )
+    assert torch.equal(module.ngram_embedding._staging[:6], table[ids])
+
+
 def test_qwen4_exp_model_state_prepares_stable_dummy_ngram_inputs() -> None:
     model_state = object.__new__(Qwen4ExpModelState)
     model_state.uses_ngram_embedding = True
@@ -287,3 +326,107 @@ def test_qwen4_exp_model_state_prepares_stable_dummy_ngram_inputs() -> None:
     )
     assert second["query_start_loc"].data_ptr() == query_start_loc_ptr
     assert second["ngram_context"].data_ptr() == ngram_context_ptr
+
+
+@contextmanager
+def _captured_transformers_warnings() -> Iterator[list[str]]:
+    """Records what transformers logs, without relying on propagation.
+
+    transformers keeps its own logger hierarchy, and the rope validator reaches
+    it through ``warning_once``, so a root-level capture can miss the record
+    entirely.
+    """
+    records: list[str] = []
+
+    class _Collect(logging.Handler):
+        def emit(self, record: logging.LogRecord) -> None:
+            records.append(record.getMessage())
+
+    handler = _Collect()
+    logger = logging.getLogger("transformers")
+    previous_level = logger.level
+    logger.addHandler(handler)
+    logger.setLevel(logging.WARNING)
+    try:
+        yield records
+    finally:
+        logger.removeHandler(handler)
+        logger.setLevel(previous_level)
+
+
+@pytest.mark.parametrize("wrapped_config", [False, True])
+def test_rope_validation_accepts_mrope_keys_and_still_warns_for_unknown_keys(
+    wrapped_config: bool,
+) -> None:
+    """M-RoPE is driven through two rope_parameters keys transformers does not
+    know for rope_type="default". Declaring them must not silence the
+    unrecognized-key warning for every other key."""
+
+    def build(rope_parameters: dict) -> None:
+        text_config = _text_config(
+            rope_parameters=dict(rope_parameters), rope_theta=10_000.0
+        )
+        if wrapped_config:
+            Qwen4ExpConfig(
+                architectures=["Qwen4ExpForConditionalGeneration"],
+                text_config=text_config.to_dict(),
+            )
+
+    with _captured_transformers_warnings() as records:
+        build(
+            {
+                "rope_type": "default",
+                "mrope_section": [16, 24, 24],
+                "mrope_interleaved": True,
+            }
+        )
+    assert not [line for line in records if "Unrecognized keys" in line], records
+
+    # warning_once dedupes on the message, so each parametrization needs a key
+    # no earlier run has reported.
+    unknown_key = f"not_a_rope_key_{int(wrapped_config)}"
+    with _captured_transformers_warnings() as records:
+        build({"rope_type": "default", unknown_key: 1})
+    assert [line for line in records if unknown_key in line], records
+
+
+@pytest.mark.parametrize("is_last_rank", [True, False])
+def test_mtp_hidden_buffer_is_allocated_on_the_configured_device(
+    is_last_rank: bool,
+) -> None:
+    """The drafter's hidden buffer belongs on the device the config names, not
+    on whichever device happens to be current."""
+    vllm_config = SimpleNamespace(
+        speculative_config=SimpleNamespace(method="mtp"),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
+        model_config=SimpleNamespace(dtype=torch.bfloat16),
+        # "meta" is never the ambient default, so an implicit allocation lands
+        # somewhere else and this test sees it.
+        device_config=SimpleNamespace(device=torch.device("meta")),
+    )
+    config = SimpleNamespace(hc_count=2, hidden_size=4)
+
+    buffer = make_mtp_hidden_buffer(vllm_config, config, is_last_rank=is_last_rank)
+
+    if not is_last_rank:
+        assert buffer is None
+        return
+    assert buffer is not None
+    assert buffer.device.type == "meta"
+    assert buffer.shape == (8, 8)
+    assert buffer.dtype is torch.bfloat16
+
+
+def test_mtp_hidden_buffer_is_absent_without_mtp_speculation() -> None:
+    vllm_config = SimpleNamespace(
+        speculative_config=SimpleNamespace(method="eagle"),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
+        model_config=SimpleNamespace(dtype=torch.bfloat16),
+        device_config=SimpleNamespace(device=torch.device("meta")),
+    )
+    config = SimpleNamespace(hc_count=2, hidden_size=4)
+
+    assert make_mtp_hidden_buffer(vllm_config, config, is_last_rank=True) is None
+
+    vllm_config.speculative_config = None
+    assert make_mtp_hidden_buffer(vllm_config, config, is_last_rank=True) is None

tests/models/qwen4_exp/test_ple.py

@@ -2,13 +2,18 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import logging
+import mmap
+import os
 from dataclasses import dataclass
 from itertools import accumulate
+from pathlib import Path
 from types import SimpleNamespace
 
 import pytest
 import torch
 from compressed_tensors.quantization import QuantizationArgs
+from safetensors import safe_open
+from safetensors.torch import load, load_file, save_file
 from torch import nn
 from torch.nn import functional as F
 from transformers import Qwen4ExpTextConfig
@@ -16,6 +21,7 @@ from transformers import Qwen4ExpTextConfig
 import vllm.model_executor.layers.vocab_parallel_embedding as embedding_module
 import vllm.model_executor.parameter as parameter_module
 import vllm.models.qwen4_exp.common.ngram_embedding as ngram_embedding_module
+import vllm.models.qwen4_exp.nvidia.ngram_embedding as nvidia_ngram_embedding_module
 from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
 from vllm.config import SpeculativeConfig, VllmConfig
 from vllm.config.compilation import CUDAGraphMode
@@ -44,6 +50,7 @@ from vllm.models.qwen4_exp.nvidia.ngram_embedding import (
     Qwen4ExpNGramEmbedding,
     Qwen4ExpPLEDeviceEmbedding,
     Qwen4ExpPLEEmbeddingMethod,
+    Qwen4ExpPLEFileGatherEmbedding,
     Qwen4ExpPLEFp8EmbeddingMethod,
     Qwen4ExpPLENvFp4EmbeddingMethod,
     Qwen4ExpPLEPinnedHostEmbedding,
@@ -1106,6 +1113,351 @@ def test_ple_fp8_embedding_supports_mixed_precision_config(
     )
 
 
+def _patch_ple_construction(
+    monkeypatch: pytest.MonkeyPatch,
+    *,
+    engram_config: SimpleNamespace | None = None,
+    load_format: str = "auto",
+) -> None:
+    _mock_etp_group(monkeypatch)
+    for module in (embedding_module, parameter_module):
+        monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 0)
+        monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 1)
+    vllm_config = SimpleNamespace(
+        load_config=SimpleNamespace(load_format=load_format),
+        engram_config=engram_config,
+    )
+    monkeypatch.setattr(
+        nvidia_ngram_embedding_module, "get_current_vllm_config", lambda: vllm_config
+    )
+
+
+def _make_file_gather_embedding(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+    *,
+    fp8: bool = False,
+    load_format: str = "auto",
+    bind: bool = False,
+) -> tuple[Qwen4ExpNGramEmbedding, torch.Tensor]:
+    """Build a two-head, 12-token host-file-gather embedding and its saved table.
+
+    The 21-row table is saved to ``tmp_path / "model.safetensors"`` as four
+    checkpoint shards of stride 6, so the last shard holds 3 rows.
+    """
+    _patch_ple_construction(monkeypatch, load_format=load_format)
+    method = (
+        Qwen4ExpPLEFp8EmbeddingMethod()
+        if fp8
+        else Qwen4ExpPLEUnquantizedEmbeddingMethod()
+    )
+    module = Qwen4ExpNGramEmbedding.__new__(Qwen4ExpNGramEmbedding)
+    nn.Module.__init__(module)
+    module.use_sequence_parallel = False
+    module.split_ngram_parts = 4
+    module.eos_token_id = 0
+    module.ngram_size = 2
+    module.heads_per_ngram = 2
+    module.register_buffer("layer_multipliers", torch.tensor([3, 5]))
+    module.register_buffer("ngram_heads_vocab_sizes", torch.tensor([13, 7]))
+    module.register_buffer("ngram_heads_offsets", torch.tensor([0, 13]))
+    module.ngram_embedding = Qwen4ExpPLEFileGatherEmbedding(
+        21,
+        2,
+        params_dtype=torch.bfloat16,
+        padding_size=1,
+        prefix="test.ple_embedding",
+        embedding_method=method,
+        num_ngram_heads=2,
+        max_total_tokens=12,
+    )
+    dtype = torch.float8_e4m3fn if fp8 else torch.bfloat16
+    table = torch.arange(42, dtype=torch.float32).reshape(21, 2).to(dtype)
+    checkpoint = {
+        f"ngram_embedding.shard_{index}.weight": shard.clone()
+        for index, shard in enumerate(table.split(6))
+    }
+    if fp8:
+        checkpoint["ngram_embedding.weight_scale"] = torch.tensor([0.25]).bfloat16()
+    save_file(checkpoint, tmp_path / "model.safetensors")
+    if bind:
+        _load_file_shards(module, tmp_path / "model.safetensors")
+        module.ngram_embedding.bind_file_shards()
+    return module, table
+
+
+def _load_file_shards(module: Qwen4ExpNGramEmbedding, path: Path) -> set[str]:
+    """Load ``path`` through ``load_weights`` as the lazy safetensors iterator does."""
+    with safe_open(path, framework="pt") as checkpoint:
+        return module.load_weights(
+            (name, checkpoint.get_tensor(name))
+            for name in checkpoint.keys()  # noqa: SIM118
+        )
+
+
+def _dev_zero_shards() -> list[tuple[str, torch.Tensor]]:
+    """Return shards in a ``/dev/zero`` mapping, which is how pinned memory maps."""
+    with open("/dev/zero", "r+b") as dev_zero:
+        mapping = mmap.mmap(dev_zero.fileno(), mmap.PAGESIZE)
+    table = torch.frombuffer(mapping, dtype=torch.uint8)[:84].view(torch.bfloat16)
+    return [
+        (f"ngram_embedding.shard_{index}.weight", shard)
+        for index, shard in enumerate(table.reshape(21, 2).split(6))
+    ]
+
+
+@pytest.mark.parametrize(
+    ("engram_config", "backend", "rows"),
+    [
+        (
+            SimpleNamespace(host_file_gather=True, cpu_offload=True),
+            Qwen4ExpPLEFileGatherEmbedding,
+            0,
+        ),
+        (
+            SimpleNamespace(host_file_gather=False, cpu_offload=False),
+            Qwen4ExpPLEDeviceEmbedding,
+            24,
+        ),
+        (None, Qwen4ExpPLEDeviceEmbedding, 24),
+    ],
+)
+def test_ngram_embedding_selects_backend_from_engram_config(
+    monkeypatch, engram_config, backend, rows
+) -> None:
+    """File gather wins with no rows; ``cpu_offload=False`` keeps rows on device."""
+    _patch_ple_construction(monkeypatch, engram_config=engram_config)
+    config = SimpleNamespace(
+        ngram_size=2,
+        heads_per_ngram=2,
+        eos_token_id=0,
+        vocab_size=64,
+        ngram_vocab_size_base=8,
+        make_ngram_vocab_size_divisible_by=1,
+    )
+
+    module = Qwen4ExpNGramEmbedding(config, 4, 0, 4, data_parallel_rank=0, prefix="t")
+
+    assert tuple(module.ngram_embedding.weight.shape) == (rows, 2)
+    assert isinstance(module.ngram_embedding, backend)
+
+
+def test_ngram_embedding_default_cpu_offload_selects_pinned_backend(
+    monkeypatch,
+) -> None:
+    """Default ``cpu_offload=True`` selects pinned storage and requires UVA."""
+    _patch_ple_construction(
+        monkeypatch,
+        engram_config=SimpleNamespace(host_file_gather=False, cpu_offload=True),
+    )
+    monkeypatch.setattr(ngram_embedding_module, "is_uva_available", lambda: False)
+    config = SimpleNamespace(
+        ngram_size=2,
+        heads_per_ngram=2,
+        eos_token_id=0,
+        vocab_size=64,
+        ngram_vocab_size_base=8,
+        make_ngram_vocab_size_divisible_by=1,
+    )
+
+    with pytest.raises(RuntimeError, match="CPU offload requires UVA"):
+        Qwen4ExpNGramEmbedding(config, 4, 0, 4, data_parallel_rank=0, prefix="t")
+
+
+@pytest.mark.parametrize(
+    ("read_shards", "error"),
+    [
+        pytest.param(
+            lambda path: [("ngram_embedding.shard_0.weight", torch.zeros(6, 2))],
+            "dtype torch.float32 must match the embedding dtype torch.bfloat16",
+            id="wrong_dtype",
+        ),
+        pytest.param(
+            lambda path: [*load_file(path).items()] * 2,
+            "Duplicate PLE embedding shard 0",
+            id="duplicate_shard",
+        ),
+        pytest.param(
+            lambda path: [s for s in load_file(path).items() if "_2." not in s[0]],
+            "PLE shards cover 15 of 21 rows",
+            id="missing_shard",
+        ),
+        pytest.param(lambda path: [], "PLE shards cover 0 of 21 rows", id="no_shards"),
+        pytest.param(
+            lambda path: [(n, t.clone()) for n, t in load_file(path).items()],
+            "not a file-backed view",
+            id="anonymous",
+        ),
+        pytest.param(
+            lambda path: load(path.read_bytes()).items(),
+            "not a file-backed view",
+            id="heap",
+        ),
+        pytest.param(
+            lambda path: _dev_zero_shards(), "not a file-backed view", id="dev_zero"
+        ),
+    ],
+)
+def test_file_gather_rejects_shards_that_cannot_serve_every_row_from_a_file(
+    monkeypatch, tmp_path, read_shards, error
+) -> None:
+    """A real load must fail instead of serving zeros or a copied table."""
+    module, _ = _make_file_gather_embedding(monkeypatch, tmp_path)
+
+    with pytest.raises(ValueError, match=error):
+        module.load_weights(read_shards(tmp_path / "model.safetensors"))
+        module.ngram_embedding.bind_file_shards()
+
+
+@pytest.mark.parametrize("fp8", [False, True])
+def test_file_gather_binds_zero_copy_shards_through_loader_hook(
+    monkeypatch, tmp_path, fp8
+) -> None:
+    """Fails if safetensors stops returning zero-copy views of the file."""
+    module, table = _make_file_gather_embedding(monkeypatch, tmp_path, fp8=fp8)
+    embedding = module.ngram_embedding
+
+    loaded = _load_file_shards(module, tmp_path / "model.safetensors")
+    embedding.quant_method.process_weights_after_loading(embedding)
+    embedding.quant_method.process_weights_after_loading(embedding)
+
+    assert loaded == {"ngram_embedding.weight"} | (
+        {"ngram_embedding.weight_scale"} if fp8 else set()
+    )
+    assert embedding._bound
+    stored = torch.cat([embedding._shards[index] for index in range(4)])
+    assert torch.equal(stored, table.view(torch.uint8))
+
+
+def test_file_gather_binds_shards_opened_through_a_symlink(
+    monkeypatch, tmp_path
+) -> None:
+    """HF cache snapshots link to suffix-less blobs, which is the mapped path."""
+    module, _ = _make_file_gather_embedding(monkeypatch, tmp_path)
+    blob = (tmp_path / "model.safetensors").rename(tmp_path / "blob")
+    snapshot = tmp_path / "snapshot.safetensors"
+    snapshot.symlink_to(blob)
+    _load_file_shards(module, snapshot)
+
+    module.ngram_embedding.bind_file_shards()
+
+    assert module.ngram_embedding._bound
+
+
+def test_file_gather_rejects_reload_and_stays_poisoned(monkeypatch, tmp_path) -> None:
+    module, table = _make_file_gather_embedding(monkeypatch, tmp_path, bind=True)
+
+    with pytest.raises(RuntimeError, match="does not support weight reload"):
+        module.load_weights([("ngram_embedding.shard_0.weight", table[:6])])
+    with pytest.raises(RuntimeError, match="unbound or was reloaded"):
+        module.ngram_embedding.stage_rows(1)
+
+
+@pytest.mark.parametrize(
+    ("bad_id", "error"),
+    [(21, IndexError), (24, IndexError), (-1, RuntimeError)],
+    ids=["past_last_shard_rows", "past_last_shard", "negative"],
+)
+def test_file_gather_out_of_range_id_keeps_previous_staged_rows(
+    monkeypatch, tmp_path, bad_id, error
+) -> None:
+    """Out-of-range ids fail before the H2D copy, so no wrong row is served."""
+    module, table = _make_file_gather_embedding(monkeypatch, tmp_path, bind=True)
+    embedding = module.ngram_embedding
+    ids = torch.tensor(
+        [
+            [20, 0],
+            [19, 1],
+            [18, 2],
+            [17, 3],
+            [16, 4],
+            [15, 5],
+            [14, 6],
+            [13, 7],
+            [12, 8],
+            [11, 9],
+            [10, 10],
+        ]
+    )
+    embedding._host_ids[:11] = ids
+    embedding.stage_rows(11)
+    embedding._host_ids[5, 1] = bad_id
+
+    with pytest.raises(error):
+        embedding.stage_rows(11)
+
+    assert torch.equal(embedding._staging[:11], table[ids])
+
+
+def test_file_gather_reads_ahead_the_coalesced_pages_of_staged_rows(
+    monkeypatch, tmp_path
+) -> None:
+    """Rows start 34 bytes before page 2 of a file mapped from page 1.
+
+    Rows 0-7 are in page 1, row 8 straddles pages 1 and 2, and the rows of shard 1
+    arrive unsorted.
+    """
+    module, table = _make_file_gather_embedding(monkeypatch, tmp_path)
+    embedding, page = module.ngram_embedding, mmap.PAGESIZE
+    start = 2 * page - 34
+    path = tmp_path / "table.bin"
+    path.write_bytes(bytes(start) + table.view(torch.uint8).numpy().tobytes())
+    with open(path, "rb") as file:
+        mapping = mmap.mmap(file.fileno(), 0, offset=page, access=mmap.ACCESS_COPY)
+    rows = torch.frombuffer(mapping, dtype=torch.uint8)[start - page :]
+    module.load_weights(
+        (f"ngram_embedding.shard_{index}.weight", shard)
+        for index, shard in enumerate(rows.view(table.dtype).reshape(21, 2).split(6))
+    )
+    embedding.bind_file_shards()
+    calls = []
+    monkeypatch.setattr(
+        nvidia_ngram_embedding_module.os,
+        "posix_fadvise",
+        lambda fd, *args: calls.append((os.fstat(fd).st_ino, *args)),
+    )
+    ids = torch.tensor([[0, 1], [8, 7]])
+    embedding._host_ids[:2] = ids
+
+    embedding.stage_rows(2)
+
+    inode, willneed = path.stat().st_ino, os.POSIX_FADV_WILLNEED
+    assert calls == [(inode, page, page, willneed), (inode, page, 2 * page, willneed)]
+    assert torch.equal(embedding._staging[:2], table[ids])
+
+
+def test_file_gather_dummy_load_stages_zeros(monkeypatch, tmp_path) -> None:
+    module, _ = _make_file_gather_embedding(monkeypatch, tmp_path, load_format="dummy")
+    embedding = module.ngram_embedding
+    embedding.quant_method.process_weights_after_loading(embedding)
+    embedding._host_rows.view(torch.uint8).fill_(0xFF)
+    embedding._staging.view(torch.uint8).fill_(0xFF)
+
+    embedding.stage_rows(4)
+
+    staged = embedding._staging.view(torch.uint8)
+    assert torch.all(staged[:4] == 0)
+    assert torch.all(staged[4:] == 0xFF)
+
+
+def test_file_gather_forward_reads_staging_without_host_work(
+    monkeypatch, tmp_path
+) -> None:
+    """The forward must be a fixed-address read so FULL CUDA graphs replay it."""
+    module, _ = _make_file_gather_embedding(monkeypatch, tmp_path)
+    staging = module.ngram_embedding._staging
+    staging.view(torch.uint8).fill_(0xFF)
+    step = (torch.arange(6), torch.tensor([0, 6]), torch.tensor([[1]]))
+    hidden_states = torch.zeros(6, 8, dtype=torch.bfloat16)
+
+    assert module.start_prefetch(*step) is None
+    output = module(hidden_states, *step)
+
+    assert torch.all(staging.view(torch.uint8) == 0xFF)
+    assert output.data_ptr() == staging.data_ptr()
+    assert output.shape == (6, 4)
+
+
 def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None:
     conv_state_len = 6
     dilation = 2

tests/models/qwen4_exp/test_qsa_reference.py

@@ -506,6 +506,67 @@ def test_qsa_compressed_metadata_keeps_dummy_slots_inert() -> None:
 
 
 @requires_qsa_kernels
+@requires_qsa_kernels
+@pytest.mark.parametrize("layout", ["plain", "compressed", "circular"])
+def test_qsa_draft_decode_metadata_update_matches_rebuild(layout: str) -> None:
+    """A fused draft step's in-place update equals a fresh build at seq_lens + 1."""
+    device = torch.device("cuda")
+    compress_ratio = 4 if layout != "plain" else 1
+
+    def make_builder() -> QSAMetadataBuilder:
+        builder = QSAMetadataBuilder.__new__(QSAMetadataBuilder)
+        builder.compress_ratio = compress_ratio
+        builder.reorder_batch_threshold = 1
+        builder.is_circular_buffer = layout == "circular"
+        builder.kv_cache_spec = SimpleNamespace(block_size=4)
+        builder.storage_block_size = 16 if layout == "compressed" else 4
+        builder.token_to_req_buffer = torch.empty(8, dtype=torch.int32, device=device)
+        builder.slot_mapping_buffer = torch.empty(8, dtype=torch.int64, device=device)
+        builder.logical_positions_buffer = torch.empty(
+            8, dtype=torch.int64, device=device
+        )
+        builder.visible_blocks_buffer = torch.empty(8, dtype=torch.int32, device=device)
+        builder.request_capacity = 8
+        builder.k_work_metadata_buffer = torch.empty(
+            8 if layout == "compressed" else 0, 2, dtype=torch.int32, device=device
+        )
+        return builder
+
+    query_start_loc = torch.tensor([0, 1, 2, 3], dtype=torch.int32, device=device)
+    token_to_req = torch.tensor([0, 1, 2], dtype=torch.int32, device=device)
+
+    def make_common(seq_lens: torch.Tensor) -> SimpleNamespace:
+        return SimpleNamespace(
+            num_actual_tokens=3,
+            num_reqs=3,
+            max_query_len=1,
+            max_seq_len=int(seq_lens.max()),
+            query_start_loc=query_start_loc,
+            query_start_loc_cpu=query_start_loc.cpu(),
+            seq_lens=seq_lens,
+            slot_mapping=torch.tensor([5, 9, 20], dtype=torch.int64, device=device),
+            block_table_tensor=torch.tensor(
+                [[3, 1], [0, 4], [2, 5]], dtype=torch.int32, device=device
+            ),
+            token_to_req_indices=lambda buffer: buffer.copy_(token_to_req),
+        )
+
+    # 7 -> 8 crosses a compression boundary, so visible_blocks must change.
+    seq_lens = torch.tensor([7, 11, 16], dtype=torch.int32, device=device)
+    builder = make_builder()
+    metadata = builder.build(0, make_common(seq_lens))
+    positions_before = metadata.logical_positions.clone()
+
+    seq_lens.add_(1)
+    builder.update_draft_decode_metadata(metadata)
+    expected = make_builder().build(0, make_common(seq_lens.clone()))
+
+    assert not torch.equal(metadata.logical_positions, positions_before)
+    for name in ("token_to_req", "logical_positions", "visible_blocks", "slot_mapping"):
+        torch.testing.assert_close(getattr(metadata, name), getattr(expected, name))
+    torch.testing.assert_close(metadata.k_work_metadata, expected.k_work_metadata)
+
+
 @pytest.mark.usefixtures("default_vllm_config")
 def test_qsa_unfused_cache_update_ignores_padded_qk() -> None:
     """Padded projected Q/K rows must not affect either side cache."""

tests/models/test_abliteration.py

@@ -0,0 +1,280 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the runtime refusal-direction projection."""
+
+import json
+import math
+
+import pytest
+import torch
+
+from vllm import envs
+from vllm.model_executor.models import abliteration
+
+HIDDEN = 8
+
+
+@pytest.fixture(autouse=True)
+def clear_direction_cache():
+    abliteration._read.cache_clear()
+    abliteration.clear_logs()
+    yield
+    abliteration._read.cache_clear()
+    abliteration.clear_logs()
+
+
+def write_direction(tmp_path, name="r.json", size=HIDDEN):
+    direction = [math.sin(i) for i in range(size)]
+    path = tmp_path / name
+    path.write_text(json.dumps({"hidden_size": size, "refusal_direction": direction}))
+    return str(path)
+
+
+def registered(module, size=HIDDEN, dtype=torch.float32):
+    abliteration.register(module, size, dtype)
+    return module
+
+
+def test_off_is_a_noop_and_still_registers(tmp_path, monkeypatch):
+    """With the switch off the buffer exists as zeros and x is unchanged."""
+    monkeypatch.delenv("VLLM_ABLITERATION", raising=False)
+    monkeypatch.delenv("QWEN38_ABLITERATION", raising=False)
+    module = registered(torch.nn.Module())
+    assert module._abliteration_on is False
+    assert torch.count_nonzero(module._abliteration_r) == 0
+
+    x = torch.randn(3, HIDDEN)
+    expected = x.clone()
+    assert abliteration.project_(module, x) is x
+    abliteration.project_(module, x)
+    assert torch.equal(x, expected)
+
+
+def test_projection_drops_the_component(tmp_path, monkeypatch):
+    """After projection the tensor keeps no component along the direction."""
+    monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path))
+    module = registered(torch.nn.Module())
+    assert module._abliteration_on is True
+
+    r = module._abliteration_r
+    x = torch.randn(4, HIDDEN)
+    abliteration.project_(module, x)
+    assert torch.allclose(x @ r, torch.zeros(4), atol=1e-5)
+    assert not torch.allclose(x, torch.zeros_like(x))
+
+    # Projecting an already projected tensor changes nothing.
+    projected = x.clone()
+    abliteration.project_(module, x)
+    assert torch.allclose(x, projected, atol=1e-6)
+
+
+def test_registered_direction_is_unit(tmp_path, monkeypatch):
+    monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path))
+    module = registered(torch.nn.Module())
+    assert torch.allclose(module._abliteration_r.norm(), torch.tensor(1.0), atol=1e-6)
+
+
+def test_non_contiguous_input_is_projected(tmp_path, monkeypatch):
+    monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path))
+    module = registered(torch.nn.Module())
+    x = torch.randn(4, HIDDEN * 2)[:, ::2]
+    assert not x.is_contiguous()
+    abliteration.project_(module, x)
+    assert torch.allclose(x @ module._abliteration_r, torch.zeros(4), atol=1e-5)
+
+
+def test_width_mismatch_is_skipped(tmp_path, monkeypatch):
+    """A tensor without the full hidden dim must not be projected."""
+    monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path))
+    module = registered(torch.nn.Module())
+    x = torch.randn(2, HIDDEN // 2)
+    before = x.clone()
+    abliteration.project_(module, x)
+    assert torch.equal(x, before)
+
+
+def test_hidden_size_mismatch_raises(tmp_path, monkeypatch):
+    """A direction for another model fails at boot, not silently."""
+    monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path, size=HIDDEN))
+    with pytest.raises(ValueError, match="holds 8 values"):
+        registered(torch.nn.Module(), size=HIDDEN + 1)
+
+
+def test_first_projection_logs_that_it_is_active(tmp_path, monkeypatch, caplog):
+    """The boot log proves the switch is live, not just loaded."""
+    import logging
+
+    monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path))
+    module = registered(torch.nn.Linear(HIDDEN, HIDDEN))
+    with caplog.at_level(
+        logging.INFO, logger="vllm.model_executor.models.abliteration"
+    ):
+        abliteration.project_(module, torch.randn(2, HIDDEN))
+    assert "active, projected a 8-wide tensor" in caplog.text
+
+
+def test_flash_next_direction_is_rejected_by_the_27b_model(tmp_path, monkeypatch):
+    """Flash-Next stores 2560 values; the dense 27B stream is 5120 wide."""
+    path = write_direction(tmp_path, name="flash_next.json", size=2560)
+    monkeypatch.setenv("VLLM_ABLITERATION", path)
+    with pytest.raises(
+        ValueError, match="holds 2560 values; the model's hidden size is 5120"
+    ):
+        registered(torch.nn.Module(), size=5120)
+
+
+def test_half_width_activation_is_skipped(tmp_path, monkeypatch):
+    """A 2560-wide tensor cannot be projected by the 5120 direction."""
+    monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path, size=5120))
+    module = registered(torch.nn.Module(), size=5120)
+    x = torch.randn(4, 2560)
+    before = x.clone()
+    abliteration.project_(module, x)
+    assert torch.equal(x, before)
+
+
+def test_zero_direction_raises(tmp_path, monkeypatch):
+    path = tmp_path / "zero.json"
+    path.write_text(json.dumps({"refusal_direction": [0.0] * HIDDEN}))
+    monkeypatch.setenv("VLLM_ABLITERATION", str(path))
+    with pytest.raises(ValueError, match="zero or non-finite"):
+        abliteration.load_direction(HIDDEN)
+
+
+def test_bare_name_resolves_in_builtin_dir(tmp_path, monkeypatch):
+    path = write_direction(tmp_path, name="Uncensored.json")
+    monkeypatch.setattr(abliteration, "BUILTIN_DIRS", (str(tmp_path),))
+    monkeypatch.setenv("VLLM_ABLITERATION", "Uncensored")
+    assert abliteration.direction_path() == path
+
+
+def test_missing_direction_file_raises(tmp_path, monkeypatch):
+    monkeypatch.setattr(abliteration, "BUILTIN_DIRS", (str(tmp_path / "nope"),))
+    monkeypatch.setenv("VLLM_ABLITERATION", "Uncensored")
+    with pytest.raises(FileNotFoundError):
+        abliteration.load_direction(HIDDEN)
+
+
+def test_meta_device_turns_projection_off(tmp_path, monkeypatch):
+    """A meta-initialized module cannot keep the direction, so it stays off."""
+    monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path))
+    module = torch.nn.Module()
+    with torch.device("meta"):
+        abliteration.register(module, HIDDEN, torch.float32)
+    assert module._abliteration_r.device.type == "meta"
+    assert module._abliteration_on is False
+
+
+def test_legacy_env_name_is_accepted(tmp_path, monkeypatch):
+    """QWEN38_ABLITERATION still works, as existing deployments set it."""
+    path = write_direction(tmp_path)
+    monkeypatch.delenv("VLLM_ABLITERATION", raising=False)
+    monkeypatch.setenv("QWEN38_ABLITERATION", path)
+    assert path == envs.VLLM_ABLITERATION
+    assert abliteration.direction_path() == path
+
+
+def test_compile_projects_without_a_graph_break(tmp_path, monkeypatch):
+    """The projection stays traceable, so it does not split the compiled model."""
+    monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path))
+
+    class Project(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            abliteration.register(self, HIDDEN, torch.float32)
+
+        def forward(self, x):
+            return abliteration.project_(self, x)
+
+    module = Project().eval()
+    x = torch.randn(16, HIDDEN)
+    expected = abliteration.project_(module, x.clone())
+    got = torch.compile(module, fullgraph=True)(x.clone())
+    assert torch.allclose(got, expected, atol=1e-6)
+
+
+def qwen3_5_config(dense: bool):
+    """A two-layer Qwen3.5 config, one full-attention and one linear layer."""
+    if dense:
+        from transformers import Qwen3_5TextConfig as config_cls
+
+        extra = {"intermediate_size": 256, "hidden_act": "silu"}
+    else:
+        from transformers import Qwen3_5MoeTextConfig as config_cls
+
+        extra = {
+            "moe_intermediate_size": 32,
+            "shared_expert_intermediate_size": 32,
+            "num_experts": 2,
+            "num_experts_per_tok": 1,
+            "n_group": 1,
+            "topk_group": 1,
+            "norm_topk_prob": False,
+            "decoder_sparse_step": 1,
+            "mlp_only_layers": [],
+        }
+    return config_cls(
+        vocab_size=64,
+        hidden_size=HIDDEN,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=32,
+        rms_norm_eps=1e-6,
+        layer_types=["full_attention", "linear_attention"],
+        linear_key_head_dim=32,
+        linear_value_head_dim=32,
+        linear_num_key_heads=2,
+        linear_num_value_heads=2,
+        num_nextn_predict_layers=1,
+        **extra,
+    )
+
+
+@pytest.mark.parametrize("dense", [True, False])
+def test_qwen3_5_model_projects_embeddings(tmp_path, monkeypatch, dense, dist_init):
+    """A real Qwen3.5 model registers the direction and drops it from embeddings."""
+    from types import SimpleNamespace
+
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.model_executor.models.qwen3_5 import Qwen3_5Model
+
+    config = qwen3_5_config(dense)
+    path = write_direction(tmp_path, size=HIDDEN)
+    monkeypatch.setenv("VLLM_ABLITERATION", path)
+
+    vllm_config = VllmConfig()
+    vllm_config.model_config = SimpleNamespace(
+        hf_config=config,
+        hf_text_config=config,
+        dtype=torch.bfloat16,
+        max_model_len=128,
+        is_diffusion=False,
+        head_dtype=None,
+        is_mm_prefix_lm=False,
+        multimodal_config=None,
+        quantization_config=None,
+        runner_type="generate",
+        is_moe=not dense,
+        logits_processors=None,
+        rswa_window=None,
+    )
+    vllm_config.quant_config = None
+
+    with set_current_vllm_config(vllm_config):
+        model = Qwen3_5Model(vllm_config=vllm_config)
+    with torch.no_grad():
+        for param in model.parameters():
+            param.data.normal_(0, 0.02)
+
+    assert model._abliteration_on is True
+    for layer in model.layers:
+        assert layer._abliteration_on is True
+
+    r = model._abliteration_r.float()
+    # The buffer carries the model dtype, so bf16 rounding keeps it near unit norm.
+    assert torch.allclose(r.norm(), torch.tensor(1.0), atol=2e-2)
+
+    embeddings = model.embed_input_ids(torch.arange(6)).float()
+    residual = float((embeddings @ r).abs().max())
+    scale = float(embeddings.abs().max())
+    assert residual < 1e-2 * scale

tests/models/test_deepseek_v41_decoder_replay_layers.py

@@ -165,7 +165,7 @@ def test_replay_graph_matches_eager(monkeypatch, moe_backend, ids_dtype):
     # Mocking it true would capture into the cuMem pool, which breaks on ROCm.
     cfg.use_cumem_cudagraph_pool = False
     cfg.scheduler_config = MagicMock(max_num_seqs=2, max_num_batched_tokens=64)
-    cfg.cache_config.use_kda_recoverssm = False
+    cfg.cache_config.use_recoverssm = False
     cfg.model_config.hf_config = MagicMock(hc_mult=HC, hidden_size=HIDDEN)
     cfg.model_config.dtype = torch.float32
     cfg.kernel_config.moe_backend = moe_backend

tests/test_config.py

@@ -243,7 +243,7 @@ def test_kda_recoverssm_derivation_is_revalidated():
     config = SimpleNamespace(
         cache_config=SimpleNamespace(
             use_replayssm=True,
-            use_kda_recoverssm=False,
+            use_recoverssm=False,
             mamba_cache_mode="none",
             replayssm_buffer_len=16,
         ),
@@ -263,7 +263,7 @@ def test_kda_recoverssm_derivation_is_revalidated():
 
     VllmConfig.validate_mamba_cached_kernel(config)
     assert config.cache_config.use_replayssm
-    assert config.cache_config.use_kda_recoverssm
+    assert config.cache_config.use_recoverssm
 
     config.cache_config.mamba_cache_mode = "align"
     VllmConfig.validate_mamba_cached_kernel(config)
@@ -276,7 +276,7 @@ def test_kda_recoverssm_derivation_is_revalidated():
     config.model_config.architecture = "NemotronHForCausalLM"
     config.mamba_config.backend = MambaBackendEnum.FLASHINFER
     VllmConfig.validate_mamba_cached_kernel(config)
-    assert not config.cache_config.use_kda_recoverssm
+    assert not config.cache_config.use_recoverssm
 
     config.model_config.architecture = "KimiLinearForCausalLM"
     config.parallel_config.pipeline_parallel_size = 2
@@ -1590,7 +1590,9 @@ def test_engram_tensor_parallel_size(dp_size: int, across_dp: bool, expected: in
     assert config.get_parallel_size(parallel) == expected
 
 
-@pytest.mark.parametrize("option", ["embedding_across_dp", "dp_shared_memory"])
+@pytest.mark.parametrize(
+    "option", ["embedding_across_dp", "dp_shared_memory", "host_file_gather"]
+)
 def test_engram_rejects_elastic_cross_dp(option):
     parallel = ParallelConfig(
         tensor_parallel_size=4,
@@ -1814,12 +1816,47 @@ def test_engram_draft_config_validates_target(monkeypatch, target_has_ple, expli
             assert config.engram_config.cpu_offload is True
 
 
+@pytest.mark.parametrize(
+    ("architecture", "layer_field", "rocm", "expect_error"),
+    [
+        pytest.param(
+            "DeepseekV41ForCausalLM", "engram_layer_ids", False, True, id="deepseek"
+        ),
+        pytest.param(
+            "Qwen4ExpForCausalLM", "ple_layer_ids", True, True, id="qwen4exp_rocm"
+        ),
+        pytest.param(
+            "Qwen4ExpForCausalLM", "ple_layer_ids", False, False, id="qwen4exp_cuda"
+        ),
+    ],
+)
+def test_engram_host_file_gather_requires_cuda_qwen4exp(
+    monkeypatch, architecture, layer_field, rocm, expect_error
+):
+    monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: True)
+    monkeypatch.setattr(current_platform, "is_rocm", lambda: rocm)
+    model = cast(
+        ModelConfig,
+        SimpleNamespace(
+            architecture=architecture,
+            hf_text_config=SimpleNamespace(**{layer_field: [1]}),
+        ),
+    )
+    config = EngramConfig(host_file_gather=True)
+    if expect_error:
+        with pytest.raises(ValueError, match="host_file_gather"):
+            config.verify_model_config(model)
+    else:
+        config.verify_model_config(model)
+
+
 def test_engram_hash_tracks_execution_options():
     configs = [
         EngramConfig(cpu_offload=True),
         EngramConfig(cpu_offload=False),
         EngramConfig(cpu_offload=True, embedding_across_dp=True),
         EngramConfig(cpu_offload=True, dp_shared_memory=True),
+        EngramConfig(host_file_gather=True),
     ]
     assert len({config.compute_hash() for config in configs}) == len(configs)
 

tests/v1/core/test_single_type_kv_cache_manager.py

@@ -1014,6 +1014,77 @@ def test_predictor_matches_allocator_blocks_calculation_with_admission_cap():
         total_computed = num_tokens
 
 
+def test_mamba_honors_drop_eagle_block():
+    """MambaManager must drop the EAGLE lookahead block, like FullAttentionManager.
+
+    EAGLE/MTP requires the final matched block of a cache hit be dropped and
+    recomputed: it can hold state written over draft positions that verification
+    later rejects. FullAttentionManager pops it. Before this fix MambaManager
+    accepted `drop_eagle_block` and ignored it, so its hit came back one block
+    longer, and the retained block's recurrent-state snapshot -- taken over the
+    unverified positions -- stayed hash-reachable for every later request sharing
+    the prefix (#43559).
+
+    block_size == alignment_tokens here, so the alignment loop is a no-op in both
+    managers and the eagle drop is the only behavior under test.
+    """
+    block_size = 16
+    num_blocks = 5
+    full_gid, mamba_gid = 0, 1
+
+    block_pool = BlockPool(
+        num_gpu_blocks=64, enable_caching=True, hash_block_size=block_size
+    )
+    block_hashes = [BlockHash(f"blk{i:04d}".encode()) for i in range(num_blocks)]
+
+    # Cache the same prefix chain under both groups, each on real blocks.
+    for gid in (full_gid, mamba_gid):
+        for i, (block_hash, block) in enumerate(
+            zip(block_hashes, block_pool.get_new_blocks(num_blocks))
+        ):
+            block_pool._insert_block_hash(
+                make_block_hash_with_group_id(block_hash, gid),
+                block,
+                num_tokens=(i + 1) * block_size,
+            )
+
+    full_spec = FullAttentionSpec(
+        block_size=block_size, num_kv_heads=1, head_size=1, dtype=torch.float16
+    )
+    mamba_spec = MambaSpec(
+        block_size=block_size,
+        shapes=((1,),),
+        dtypes=(torch.float16,),
+        mamba_cache_mode="align",
+    )
+
+    def hit_blocks(manager, gid, spec, drop_eagle_block):
+        return len(
+            manager.find_longest_cache_hit(
+                block_hashes=block_hashes,
+                max_length=num_blocks * block_size,
+                kv_cache_group_ids=[gid],
+                block_pool=block_pool,
+                kv_cache_spec=spec,
+                drop_eagle_block=drop_eagle_block,
+                alignment_tokens=block_size,
+            )[0][0]
+        )
+
+    # Control: with no eagle drop, both groups see the whole cached prefix.
+    assert hit_blocks(FullAttentionManager, full_gid, full_spec, False) == num_blocks
+    assert hit_blocks(MambaManager, mamba_gid, mamba_spec, False) == num_blocks
+
+    # With the eagle drop requested, both must give up the final block.
+    full_hit = hit_blocks(FullAttentionManager, full_gid, full_spec, True)
+    mamba_hit = hit_blocks(MambaManager, mamba_gid, mamba_spec, True)
+    assert full_hit == num_blocks - 1
+    assert mamba_hit == full_hit, (
+        f"MambaManager ignored drop_eagle_block: full-attention dropped the "
+        f"lookahead block (hit={full_hit}), mamba kept it (hit={mamba_hit})"
+    )
+
+
 def test_pop_blocks_for_free_trims_null_head():
     block_size = 2
     spec = SlidingWindowSpec(

tests/v1/cudagraph/test_cudagraph_manager.py

@@ -316,7 +316,7 @@ def _create_decode_vllm_config(
     num_speculative_tokens: int = 0,
     dynamic_spec_schedule: list[tuple[int, int, int]] | None = None,
     max_num_seqs: int = 8,
-    use_kda_recoverssm: bool = False,
+    use_recoverssm: bool = False,
 ) -> MagicMock:
     compilation_config = CompilationConfig(
         cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,
@@ -331,7 +331,7 @@ def _create_decode_vllm_config(
         max_num_seqs=max_num_seqs
     )
     vllm_config.parallel_config = ParallelConfig()
-    vllm_config.cache_config = SimpleNamespace(use_kda_recoverssm=use_kda_recoverssm)
+    vllm_config.cache_config = SimpleNamespace(use_recoverssm=use_recoverssm)
     vllm_config.num_speculative_tokens = num_speculative_tokens
     if dynamic_spec_schedule is None:
         vllm_config.speculative_config = None
@@ -350,7 +350,7 @@ def _make_spec_decode_manager(
     num_speculative_tokens: int = 0,
     dynamic_spec_schedule: list[tuple[int, int, int]] | None = None,
     max_num_seqs: int = 8,
-    use_kda_recoverssm: bool = False,
+    use_recoverssm: bool = False,
     varlen_decode: bool = False,
 ) -> gpu_cudagraph_utils.CudaGraphManager:
     monkeypatch.setattr(
@@ -369,7 +369,7 @@ def _make_spec_decode_manager(
             num_speculative_tokens=num_speculative_tokens,
             dynamic_spec_schedule=dynamic_spec_schedule,
             max_num_seqs=max_num_seqs,
-            use_kda_recoverssm=use_kda_recoverssm,
+            use_recoverssm=use_recoverssm,
         ),
         device=torch.device("cpu"),
         cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,

tests/v1/e2e/general/test_mamba_prefix_cache.py

@@ -1021,6 +1021,102 @@ def _run_mamba_prefix_cache_mrv1(
     cleanup_dist_env_and_memory()
 
 
+def _run_mamba_eagle_block_reuse(
+    monkeypatch: pytest.MonkeyPatch, async_scheduling: bool
+):
+    """Recreate the #43559 failure conditions: a request finishing right after a
+    block boundary it crossed with every draft rejected, then a second request
+    whose cache hit lands on that final (eagle) block.
+
+    Request A ends a few tokens past the 2nd aligned boundary; with
+    num_accepted_tokens=1 every boundary crossing happens with rejected drafts
+    in flight, and A finishes before any later step could repair its cached
+    snapshot. Request B shares the prefix: pre-fix MambaManager returns the
+    boundary block as part of the hit (B resumes on the possibly-poisoned
+    snapshot), post-fix the hit stops one block earlier and B recomputes it.
+
+    Assertions, all deterministic:
+      1. B's num_cached_tokens == BLOCK_SIZE (the eagle block was dropped).
+      2. A's cached snapshot at the boundary matches the no-spec reference.
+      3. B's state at its next boundary matches the reference -- corruption
+         resumed from the reused block would surface here (the user-visible
+         failure, at state level).
+    """
+    global async_scheduling_mode
+    async_scheduling_mode = async_scheduling
+    run_ref_mamba_state_in_subprocess()
+    apply_patch(monkeypatch)
+    full_prompt = datasets.load_dataset("heheda/a_long_article")["train"][0]["text"]
+
+    engine = LLM(
+        model=MODEL,
+        load_format="dummy",
+        enable_prefix_caching=True,
+        block_size=BLOCK_SIZE,
+        mamba_cache_mode="align",
+        speculative_config={
+            "method": "qwen3_next_mtp",
+            "num_speculative_tokens": num_speculative_tokens,
+        },
+        max_num_batched_tokens=3072,
+        hf_overrides={"num_hidden_layers": NUM_HIDDEN_LAYERS},
+        async_scheduling=async_scheduling,
+        seed=42,
+    )
+    global prompt_token_ids
+    prompt_token_ids = engine.get_tokenizer().encode(full_prompt)
+    global num_accepted_tokens
+    num_accepted_tokens = 1
+    global cur_step_action_idx, step_actions
+    cur_step_action_idx = 0
+    step_actions = []
+
+    boundary = BLOCK_SIZE * 2
+    # A decodes across BOTH boundaries (align mode only snapshots state at
+    # decode/chunk crossings, and B needs a fallback state block once the
+    # eagle block is dropped), rejecting 3 drafts every step, and finishes
+    # a few tokens past the second boundary -- no later step can repair the
+    # cached snapshot.
+    engine.generate(
+        [TokensPrompt(prompt_token_ids=prompt_token_ids[: BLOCK_SIZE - 2])],
+        SamplingParams(temperature=0.0, max_tokens=BLOCK_SIZE + 10),
+    )
+    a_boundary_state = mamba_kv_cache_dict.get(boundary)
+    assert a_boundary_state is not None, "A never cached the boundary block"
+    mamba_kv_cache_dict.clear()
+
+    # B: prefix through the boundary, prompt ending shy of the next one so a
+    # few decode steps cross it and its state gets captured for comparison.
+    next_boundary = BLOCK_SIZE * 3
+    (out_b,) = engine.generate(
+        [TokensPrompt(prompt_token_ids=prompt_token_ids[: next_boundary - 10])],
+        SamplingParams(temperature=0.0, max_tokens=15),
+    )
+    assert out_b.num_cached_tokens == BLOCK_SIZE, (
+        f"cache hit must drop the eagle block: expected {BLOCK_SIZE} cached "
+        f"tokens, got {out_b.num_cached_tokens}"
+    )
+
+    mamba_state_ref = torch.load("mamba_kv_cache_dict_ref.pth")
+    check_mamba_state_equal(mamba_state_ref, {boundary: a_boundary_state}, [boundary])
+    check_mamba_state_equal(mamba_state_ref, mamba_kv_cache_dict, [next_boundary])
+
+    mamba_kv_cache_dict.clear()
+    del engine
+    torch.accelerator.empty_cache()
+    cleanup_dist_env_and_memory()
+
+
+@create_new_process_for_each_test("spawn")
+def test_mamba_eagle_block_reuse_mrv1(monkeypatch: pytest.MonkeyPatch):
+    _run_mamba_eagle_block_reuse(monkeypatch, async_scheduling=False)
+
+
+@create_new_process_for_each_test("spawn")
+def test_mamba_eagle_block_reuse_mrv1_async(monkeypatch: pytest.MonkeyPatch):
+    _run_mamba_eagle_block_reuse(monkeypatch, async_scheduling=True)
+
+
 @create_new_process_for_each_test("spawn")
 def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch):
     _run_mamba_prefix_cache_mrv1(monkeypatch, async_scheduling=False)

tests/v1/spec_decode/test_qwen3_5_mtp_dense_head.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Dense-MTP-head detection for pack-quantized checkpoints (#58807).
+
+The whole fix hinges on reading the checkpoint's actual tensor names, so the
+tests pin the naming contract: dense means plain `mtp.*.weight` keys with no
+quantization artifact suffix anywhere under `mtp.`.
+"""
+
+import json
+from types import SimpleNamespace
+
+import pytest
+
+from vllm.model_executor.models.qwen3_5_mtp import (
+    _checkpoint_ships_dense_mtp,
+    _mtp_weight_names,
+)
+
+pytestmark = pytest.mark.cpu_test
+
+
+def _model_config_with_index(monkeypatch, names):
+    monkeypatch.setattr(
+        "vllm.model_executor.models.qwen3_5_mtp._mtp_weight_names",
+        lambda _config: names,
+    )
+    return SimpleNamespace()
+
+
+@pytest.mark.parametrize(
+    "names,expected",
+    [
+        # The reporter's checkpoint: every mtp.* tensor dense.
+        (["mtp.fc.weight", "mtp.layers.0.mlp.down_proj.weight"], True),
+        # Biases are dense too and must not count as artifacts.
+        (["mtp.fc.weight", "mtp.fc.bias"], True),
+        # A properly pack-quantized MTP head keeps the target quantization.
+        (["mtp.fc.weight_packed", "mtp.fc.weight_scale"], False),
+        # Channel-wise quantized MTP head: dense name plus scale companion.
+        (["mtp.fc.weight", "mtp.fc.weight_scale"], False),
+        # No MTP head in the checkpoint at all.
+        (["model.layers.0.self_attn.q_proj.weight"], False),
+        # Layout unknown (no index): keep current behavior.
+        (None, False),
+        ([], False),
+    ],
+)
+def test_dense_mtp_detection(monkeypatch, names, expected):
+    config = _model_config_with_index(monkeypatch, names)
+    assert _checkpoint_ships_dense_mtp(config) is expected
+
+
+def test_weight_names_read_from_local_index(tmp_path, monkeypatch):
+    index = {
+        "weight_map": {
+            "model.layers.0.self_attn.q_proj.weight_packed": "model-00001.safetensors",
+            "mtp.fc.weight": "model-mtp-bf16.safetensors",
+            "mtp.layers.0.mlp.down_proj.weight": "model-mtp-bf16.safetensors",
+        }
+    }
+    (tmp_path / "model.safetensors.index.json").write_text(json.dumps(index))
+    config = SimpleNamespace(model=str(tmp_path), revision=None)
+
+    names = _mtp_weight_names(config)
+
+    assert names == ["mtp.fc.weight", "mtp.layers.0.mlp.down_proj.weight"]
+
+
+def test_weight_names_none_without_index(tmp_path, monkeypatch):
+    monkeypatch.setattr(
+        "vllm.model_executor.models.qwen3_5_mtp.huggingface_hub.hf_hub_download",
+        lambda *args, **kwargs: (_ for _ in ()).throw(FileNotFoundError()),
+    )
+    config = SimpleNamespace(model=str(tmp_path), revision=None)
+
+    assert _mtp_weight_names(config) is None

vllm/config/cache.py

@@ -211,7 +211,7 @@ class CacheConfig:
     prefix caching yet. Mamba2 speculative decode requires FlashInfer. With
     prefix caching, flushes are most efficient when mamba_block_size is a
     multiple of replayssm_buffer_len, but this is not required."""
-    use_kda_recoverssm: bool = field(default=False, init=False)
+    use_recoverssm: bool = field(default=False, init=False)
     """Whether Kimi-K3 KDA uses RecoverSSM speculative decode."""
 
     # Will be set after profiling.

vllm/config/compilation.py

@@ -1428,10 +1428,15 @@ class CompilationConfig:
                 f"with {min_cg_attn_backend} backend (support: "
                 f"{min_cg_support})"
             )
-            if self.mode == CompilationMode.VLLM_COMPILE and (
+            if (
-                self.splitting_ops_contain_attention()
+                self.mode == CompilationMode.VLLM_COMPILE
-                or self.use_inductor_graph_partition
+                and (
-            ):
+                    self.splitting_ops_contain_attention()
+                    or self.use_inductor_graph_partition
+                )
+            ) or envs.VLLM_USE_BREAKABLE_CUDAGRAPH:
+                # Breakable CUDA graphs break at the attention kernels at runtime,
+                # with or without torch.compile.
                 msg += (
                     "; setting cudagraph_mode=PIECEWISE because "
                     "attention is compiled piecewise"

vllm/config/engram.py

@@ -37,6 +37,11 @@ class EngramConfig:
     cpu_offload: bool = True
     """Store embedding weights in pinned CPU memory for UVA lookup."""
 
+    host_file_gather: bool = False
+    """Serve Qwen4Exp PLE rows by gathering them on the CPU from the loader's
+    file-backed checkpoint views and staging them on the device each step,
+    instead of allocating the table; takes precedence over `cpu_offload`."""
+
     embedding_across_dp: bool = False
     """Shard embeddings across TP and all DP ranks when enabled.
     Otherwise, each DP rank has a separate TP-sharded embedding replica."""
@@ -66,6 +71,8 @@ class EngramConfig:
 
     def verify_model_config(self, model_config: "ModelConfig | None") -> None:
         """Reject Engram configuration for models without n-gram embeddings."""
+        from vllm.platforms import current_platform
+
         field = (
             _NGRAM_LAYER_FIELDS.get(model_config.architecture)
             if model_config is not None
@@ -80,6 +87,10 @@ class EngramConfig:
                 "EngramConfig requires a model with supported Engram "
                 "embeddings and non-empty n-gram layer ids."
             )
+        if self.host_file_gather and (
+            field != "ple_layer_ids" or current_platform.is_rocm()
+        ):
+            raise ValueError("host_file_gather requires a Qwen4Exp model on CUDA")
 
     def resolve_dp_shared_memory(self, parallel_config: "ParallelConfig") -> None:
         """Share host tables by default wherever the configuration permits."""
@@ -106,6 +117,8 @@ class EngramConfig:
             raise ValueError(
                 "Engram embedding_across_dp is not supported with elastic EP yet."
             )
+        if self.host_file_gather and parallel_config.enable_elastic_ep:
+            raise ValueError("host_file_gather is not supported with elastic EP.")
 
     def get_parallel_size(self, parallel_config: "ParallelConfig") -> int:
         """Derive the embedding group size from the parallel configuration."""

vllm/config/offload.py

@@ -73,6 +73,19 @@ class PrefetchOffloadConfig:
     parameters of each offloaded layer are offloaded.
     Uses segment matching: "w13_weight" matches "mlp.experts.w13_weight"
     but not "mlp.experts.w13_weight_scale".
+    Applies to decoder layers only; `offload_towers` offloads the full
+    tower regardless of this set.
+    """
+
+    offload_towers: bool = False
+    """Offload multimodal tower weights (vision/audio encoders) to pinned
+    host memory, and bulk-copy them to the GPU during each encoder forward.
+    Towers hold no VRAM between image encodes. Unlike UVA zero-copy, weights
+    are streamed with whole-tensor H2D copies instead of being re-read tile
+    by tile inside the GEMM inner loops, so encode cost stays near the
+    resident-tower baseline.
+    Works together with, or instead of, layer prefetch:
+    `offload_group_size` may be 0 when only towers are offloaded.
     """
 
 
@@ -83,7 +96,8 @@ class OffloadConfig:
     offload_backend: OffloadBackend = "auto"
     """The backend for weight offloading. Options:
     - "auto": Selects based on which sub-config has non-default values
-      (prefetch if offload_group_size > 0, uva if cpu_offload_gb > 0).
+      (prefetch if offload_group_size > 0 or offload_towers, uva if
+      cpu_offload_gb > 0).
     - "uva": UVA (Unified Virtual Addressing) zero-copy offloading.
     - "prefetch": Async prefetch with group-based layer offloading.
     """
@@ -97,7 +111,15 @@ class OffloadConfig:
     @model_validator(mode="after")
     def validate_offload_config(self) -> "OffloadConfig":
         """Validate offload configuration constraints."""
-        if self.offload_backend == "prefetch" or self.prefetch.offload_group_size > 0:
+        layer_offload = self.prefetch.offload_group_size > 0
+        if self.offload_backend == "prefetch" and not (
+            layer_offload or self.prefetch.offload_towers
+        ):
+            raise ValueError(
+                "offload_backend='prefetch' requires offload_group_size > 0"
+                " or offload_towers=True"
+            )
+        if layer_offload:
             if self.prefetch.offload_num_in_group > self.prefetch.offload_group_size:
                 raise ValueError(
                     f"offload_num_in_group ({self.prefetch.offload_num_in_group})"
@@ -114,7 +136,9 @@ class OffloadConfig:
 
         # Warn if both backends have non-default values
         uva_active = self.uva.cpu_offload_gb > 0
-        prefetch_active = self.prefetch.offload_group_size > 0
+        prefetch_active = (
+            self.prefetch.offload_group_size > 0 or self.prefetch.offload_towers
+        )
         if self.offload_backend == "uva" and prefetch_active:
             warnings.warn(
                 "Prefetch offload fields are set but offload_backend='uva'. "

vllm/config/vllm.py

@@ -71,6 +71,13 @@ else:
 
 logger = init_logger(__name__)
 
+# Architectures whose GDN (and PLE) layers run RecoverSSM under --use-replayssm.
+_QWEN4EXP_RECOVERSSM_ARCHS = (
+    "Qwen4ExpForCausalLM",
+    "Qwen4ExpForConditionalGeneration",
+    "Qwen4ExpMTP",
+)
+
 DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset(
     {
         "AliceAIForCausalLM",
@@ -3412,9 +3419,8 @@ class VllmConfig:
     @model_validator(mode="after")
     def validate_mamba_cached_kernel(self) -> "VllmConfig":
         if not self.cache_config.use_replayssm:
-            self.cache_config.use_kda_recoverssm = False
+            self.cache_config.use_recoverssm = False
             return self
-
         kda_architectures = (
             "KimiLinearForCausalLM",
             "KimiK3ForConditionalGeneration",
@@ -3423,11 +3429,15 @@ class VllmConfig:
             self.model_config is not None
             and self.model_config.architecture in kda_architectures
         )
-        self.cache_config.use_kda_recoverssm = (
+        is_qwen4exp_recoverssm_model = (
-            self.num_speculative_tokens > 0 and is_kda_model
+            self.model_config is not None
+            and self.model_config.architecture in _QWEN4EXP_RECOVERSSM_ARCHS
+        )
+        self.cache_config.use_recoverssm = self.num_speculative_tokens > 0 and (
+            is_kda_model or is_qwen4exp_recoverssm_model
         )
         use_mamba_replayssm_spec = (
-            self.num_speculative_tokens > 0 and not self.cache_config.use_kda_recoverssm
+            self.num_speculative_tokens > 0 and not self.cache_config.use_recoverssm
         )
 
         if self.model_config is not None and not self.model_config.supports_replayssm:
@@ -3435,6 +3445,17 @@ class VllmConfig:
                 "--use-replayssm is not supported for architecture "
                 f"{self.model_config.architecture!r}"
             )
+        if (
+            self.model_config is not None
+            and self.model_config.architecture in _QWEN4EXP_RECOVERSSM_ARCHS
+            and not self.cache_config.use_recoverssm
+        ):
+            # Qwen4Exp's GDN and PLE layers have no plain ReplaySSM path; with them the
+            # flag selects RecoverSSM, which needs speculative decoding.
+            raise ValueError(
+                "--use-replayssm on Qwen4Exp selects GDN RecoverSSM and requires "
+                "speculative decoding (num_speculative_tokens > 0)"
+            )
         if (
             self.mamba_config.backend == MambaBackendEnum.FLASHINFER
             and self.cache_config.replayssm_buffer_len > 16
@@ -3442,7 +3463,15 @@ class VllmConfig:
             raise ValueError(
                 "FlashInfer ReplaySSM requires --replayssm-buffer-len <= 16"
             )
-        if self.cache_config.use_kda_recoverssm:
+        if self.cache_config.use_recoverssm:
+            if self.model_config is not None and self.model_config.architecture not in (
+                "KimiLinearForCausalLM",
+                "KimiK3ForConditionalGeneration",
+                *_QWEN4EXP_RECOVERSSM_ARCHS,
+            ):
+                raise ValueError(
+                    "RecoverSSM is only supported for Kimi-K3 KDA and Qwen4Exp GDN"
+                )
             if self.mamba_config.enable_stochastic_rounding:
                 raise ValueError(
                     "RecoverSSM supports bfloat16/float32 "

vllm/engine/arg_utils.py

@@ -604,6 +604,7 @@ class EngineArgs:
     offload_num_in_group: int = PrefetchOffloadConfig.offload_num_in_group
     offload_prefetch_step: int = PrefetchOffloadConfig.offload_prefetch_step
     offload_params: set[str] = get_field(PrefetchOffloadConfig, "offload_params")
+    offload_towers: bool = PrefetchOffloadConfig.offload_towers
     gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization
     kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes
     max_num_batched_tokens: int | None = None
@@ -1470,6 +1471,9 @@ class EngineArgs:
         offload_group.add_argument(
             "--offload-params", **prefetch_kwargs["offload_params"]
         )
+        offload_group.add_argument(
+            "--offload-towers", **prefetch_kwargs["offload_towers"]
+        )
 
         # Multimodal related configs
         multimodal_kwargs = get_kwargs(MultiModalConfig)
@@ -2839,6 +2843,7 @@ class EngineArgs:
                 offload_num_in_group=self.offload_num_in_group,
                 offload_prefetch_step=self.offload_prefetch_step,
                 offload_params=self.offload_params,
+                offload_towers=self.offload_towers,
             ),
         )
 

vllm/envs.py

@@ -226,6 +226,7 @@ if TYPE_CHECKING:
     VLLM_REGEX_COMPILATION_TIMEOUT_S: int = 5
     VLLM_MSGPACK_ZERO_COPY_THRESHOLD: int = 256
     VLLM_ALLOW_INSECURE_SERIALIZATION: bool = False
+    VLLM_ABLITERATION: str = ""
     VLLM_DISABLE_REQUEST_ID_RANDOMIZATION: bool = False
     VLLM_NIXL_SIDE_CHANNEL_HOST: str = "localhost"
     VLLM_NIXL_SIDE_CHANNEL_PORT: int = 5600
@@ -1189,6 +1190,15 @@ environment_variables: dict[str, Callable[[], Any]] = {
         os.environ.get("VLLM_ALLOW_RUNTIME_LORA_UPDATING", "0").strip().lower()
         in ("1", "true")
     ),
+    # Path to a refusal direction (a JSON file with a "refusal_direction" list,
+    # or a safetensors file with that tensor) projected out of every
+    # residual-stream write of the served model, or a bare name resolved under
+    # /opt/qwen38/abliteration and /data/abliteration. Empty or "0" is off. See
+    # vllm/model_executor/models/abliteration.py. QWEN38_ABLITERATION is the
+    # older name some deployments still set.
+    "VLLM_ABLITERATION": lambda: (
+        os.getenv("VLLM_ABLITERATION") or os.getenv("QWEN38_ABLITERATION", "")
+    ),
     # We assume drivers can report p2p status correctly.
     # If the program hangs when using custom allreduce,
     # potantially caused by a bug in the driver (535 series),

vllm/model_executor/layers/mamba/abstract.py

@@ -81,11 +81,11 @@ class MambaBase(AttentionLayerBase):
             mamba_type=self.mamba_type,
             tp_replicated=self.is_kv_cache_tp_replicated,
             mamba_cache_mode=vllm_config.cache_config.mamba_cache_mode,
-            # ReplaySSM verifies the whole window off one checkpoint, so it
+            # RecoverSSM verifies the whole window off one checkpoint, so it
             # never writes the baseline's per-draft-token state slots.
             num_speculative_blocks=(
                 0
-                if vllm_config.cache_config.use_replayssm
+                if vllm_config.cache_config.use_recoverssm
                 else vllm_config.num_speculative_tokens
             ),
         )

vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py

@@ -512,10 +512,15 @@ class ChunkGatedDeltaRule(CustomOp):
 
 @PluggableLayer.register("qwen_gated_delta_net_attention")
 class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
-    def get_state_shape(
+    def _uses_gdn_recoverssm(self) -> bool:
-        self,
+        """One predicate for the replay state's shape, dtype and the attention
-    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
+        backend."""
-        return MambaStateShapeCalculator.gated_delta_net_state_shape(
+        from vllm.model_executor.layers.mamba.recoverssm_utils import uses_recoverssm
+
+        return uses_recoverssm(self.cache_config, self.num_spec)
+
+    def get_state_shape(self) -> tuple[tuple[int, ...], ...]:
+        shapes = MambaStateShapeCalculator.gated_delta_net_state_shape(
             self.tp_size,
             self.num_k_heads,
             self.num_v_heads,
@@ -524,6 +529,32 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
             self.conv_kernel_size,
             self.num_spec,
         )
+        if self._uses_gdn_recoverssm():
+            # per-token replay record [HV, spec_query_len, V + K + 1] fp32
+            shapes = (
+                *shapes,
+                (
+                    self.num_v_heads // self.tp_size,
+                    self.num_spec + 1,
+                    self.head_v_dim + self.head_k_dim + 1,
+                ),
+            )
+        return shapes
+
+    def get_state_dtype(self) -> tuple[torch.dtype, ...]:
+        dtypes = super().get_state_dtype()
+        if self._uses_gdn_recoverssm():
+            dtypes = (*dtypes, torch.float32)
+        return dtypes
+
+    def get_attn_backend(self):
+        if self._uses_gdn_recoverssm():
+            from vllm.v1.attention.backends.gdn_recoverssm import (
+                GDNRecoverSSMAttentionBackend,
+            )
+
+            return GDNRecoverSSMAttentionBackend
+        return super().get_attn_backend()
 
     def __init__(
         self,
@@ -691,6 +722,17 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
         self.enable_fused_gdn_spec_decode = (
             self.gdn_decode_kernel == "cuda" and self.speculative_config is not None
         )
+        # RecoverSSM verify replaces the per-draft snapshot path (Triton decode
+        # only)
+        self.use_gdn_recoverssm = self._uses_gdn_recoverssm()
+        if self.use_gdn_recoverssm:
+            self.gdn_decode_kernel = "triton"
+            self.enable_fused_gdn_spec_decode = False
+            self.enable_packed_recurrent_decode = False
+            logger.info_once(
+                "GDN RecoverSSM speculative verify active (spec_query_len %d)",
+                self.num_spec + 1,
+            )
         if self.speculative_config is not None:
             logger.info_once("GDN spec decode kernel: %s", self.gdn_decode_kernel)
 
@@ -742,7 +784,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
     def _fused_gdn_decode_unsupported_reason(
         self, vllm_config: VllmConfig
     ) -> str | None:
-        conv_state_dtype, recurrent_state_dtype = self.get_state_dtype()
+        conv_state_dtype, recurrent_state_dtype = self.get_state_dtype()[:2]
         if (
             self.gqa_interleaved_layout
             or self.head_k_dim != 128
@@ -1002,6 +1044,25 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
 
         return mixed_qkv_out, z_out, b_out, a_out
 
+    def _recoverssm_verify_qkv(self, mixed_qkv):
+        """Q/K/V for the RecoverSSM verify.
+
+        The verify kernel reads token-strided heads, so for a 2-D [tokens, q|k|v] conv
+        output with a contiguous last dimension it gets views instead of copies. Any
+        other layout falls back to the contiguous copies of rearrange_mixed_qkv().
+        Returns (query, key, value, is_view)."""
+        if mixed_qkv.ndim != 2 or mixed_qkv.stride(-1) != 1:
+            return (*self.rearrange_mixed_qkv(mixed_qkv), False)
+        q_dim = self.key_dim // self.tp_size
+        v_dim = self.value_dim // self.tp_size
+        num_tokens = mixed_qkv.shape[0]
+        query = mixed_qkv[:, :q_dim].view(1, num_tokens, -1, self.head_k_dim)
+        key = mixed_qkv[:, q_dim : 2 * q_dim].view(1, num_tokens, -1, self.head_k_dim)
+        value = mixed_qkv[:, 2 * q_dim : 2 * q_dim + v_dim].view(
+            1, num_tokens, -1, self.head_v_dim
+        )
+        return query, key, value, True
+
     def rearrange_mixed_qkv(self, mixed_qkv):
         """Split packed qkv into contiguous (1, seq, heads, dim) tensors.
 
@@ -1290,7 +1351,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
         dtype = qkv_or_qkvz.dtype
         num_k_heads = self.num_k_heads // self.tp_size
         num_v_heads = self.num_v_heads // self.tp_size
-        _, state_dtype = self.get_state_dtype()
+        state_dtype = self.get_state_dtype()[1]  # 3 dtypes with the replay record
 
         # All kernels use BT = chunk_size, so a single pass with T = chunk_size
         # is sufficient to populate every autotuner cache. Mirror the real
@@ -1602,7 +1663,11 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
                 ],
                 num_accepted_tokens=num_accepted_tokens,
                 query_start_loc=spec_query_start_loc,
-                max_query_len=spec_state_indices_tensor.size(-1),
+                max_query_len=(
+                    self.num_spec + 1
+                    if getattr(self, "use_gdn_recoverssm", False)
+                    else spec_state_indices_tensor.size(-1)
+                ),
                 validate_data=False,
             )
 
@@ -1640,7 +1705,17 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
         else:
             mixed_qkv_non_spec = None
 
-        query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec)
+        if (
+            spec_sequence_masks is not None
+            and mixed_qkv_spec is not None
+            and self.use_gdn_recoverssm
+        ):
+            query_spec, key_spec, value_spec, verify_views = (
+                self._recoverssm_verify_qkv(mixed_qkv_spec)
+            )
+        else:
+            verify_views = False
+            query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec)
 
         # Split mixed non-spec-decode+prefill to process independently
         split_non_spec = (
@@ -1703,7 +1778,43 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
         # 2. Recurrent attention
 
         # 2.1: Process the multi-query part
-        if spec_sequence_masks is not None:
+        if spec_sequence_masks is not None and getattr(
+            self, "use_gdn_recoverssm", False
+        ):
+            # verify the window off the checkpoint; the commit after sampling
+            # writes the state once
+            from vllm.model_executor.layers.mamba.gdn.recoverssm_gdn import (
+                gdn_recoverssm_verify,
+            )
+
+            _n = attn_metadata.num_spec_decodes
+            assert spec_query_start_loc is not None
+            assert spec_state_indices_tensor is not None
+            core_attn_out_spec = gdn_recoverssm_verify(
+                self.A_log,
+                a_spec,
+                b_spec,
+                self.dt_bias,
+                query_spec,
+                key_spec,
+                value_spec,
+                checkpoint_state=ssm_state,
+                replay_cache=self_kv_cache[2],
+                query_start_loc=spec_query_start_loc[: _n + 1],
+                state_indices=spec_state_indices_tensor[:_n, 0],
+                spec_query_len=self.num_spec + 1,
+                use_qk_l2norm_in_kernel=True,
+                # spec-only batch: write straight into the layer output
+                out=(
+                    core_attn_out[:num_actual_tokens].unsqueeze(0)
+                    if verify_views
+                    and mixed_qkv_non_spec is None
+                    and query_spec.shape[1] == num_actual_tokens
+                    else None
+                ),
+            )
+            last_recurrent_state = None
+        elif spec_sequence_masks is not None:
             core_attn_out_spec, last_recurrent_state = (
                 fused_sigmoid_gating_delta_rule_update(
                     A_log=self.A_log,
@@ -1855,7 +1966,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
                 0, non_spec_token_indx, core_attn_out_non_spec.squeeze(0)
             )
         elif spec_sequence_masks is not None:
-            core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0)
+            if core_attn_out_spec.data_ptr() != core_attn_out.data_ptr():
+                core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0)
         else:
             core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0)
 

vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py

@@ -0,0 +1,695 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""RecoverSSM for Gated DeltaNet (Qwen-style GDN): speculative verify from one
+checkpoint without per-draft state snapshots, and accepted-state recovery after
+sampling.
+
+GDN port of the Kimi-K3 KDA RecoverSSM path
+(vllm/models/kimi_k3/nvidia/ops/recoverssm.py). Differences from KDA: the gate is a
+scalar per value head (g = -exp(A_log) * softplus(a + dt_bias)), q/k heads are grouped
+over value heads (HV // H), and q/k are L2-normalised inside the kernel when requested.
+The verify step caches, per token, the delta-rule correction c_t = beta_t * (v_t - S'_t
+k_t) (V floats), the normalised key (K floats) and the log-decay g_t (1 float) in one
+fp32 "replay" record. Recovery folds the accepted prefix: S_n = exp(sum_t g_t) S_0 +
+sum_t exp(sum_{s>t} g_s) c_t k_t^T and writes the checkpoint once (plus the
+block-boundary state in align mode)."""
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+from vllm.model_executor.layers.mamba.ops.recoverssm import (
+    compact_conv_state_kernel,
+    prepare_commit_plan_kernel,
+)
+from vllm.model_executor.layers.mamba.recoverssm_utils import recoverssm_require
+from vllm.triton_utils import tl, triton
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
+
+
+def _require(cond: bool, msg: str) -> None:
+    recoverssm_require(cond, msg, "GDN RecoverSSM")
+
+
+@triton.jit
+def _gdn_recoverssm_verify_kernel(
+    q_ptr,
+    k_ptr,
+    v_ptr,
+    a_ptr,
+    b_ptr,
+    A_log_ptr,
+    dt_bias_ptr,
+    state_ptr,
+    replay_ptr,
+    out_ptr,
+    query_start_loc_ptr,
+    state_indices_ptr,
+    scale,
+    softplus_beta,
+    softplus_threshold,
+    null_block_id,
+    stride_q_token,
+    stride_k_token,
+    stride_v_token,
+    stride_a_token,
+    stride_b_token,
+    stride_state_block,
+    stride_state_head,
+    stride_state_v,
+    stride_state_k,
+    stride_replay_block,
+    stride_replay_head,
+    stride_replay_pos,
+    stride_replay_dim,
+    stride_out_token,
+    stride_qsl,
+    stride_si,
+    H: tl.constexpr,
+    HV: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    SPEC_QUERY_LEN: tl.constexpr,
+    USE_QK_L2NORM: tl.constexpr,
+):
+    pid_v = tl.program_id(0)
+    pid_b = tl.program_id(1)
+    pid_hv = tl.program_id(2)
+    pid_h = pid_hv // (HV // H)
+    bos = tl.load(query_start_loc_ptr + pid_b * stride_qsl).to(tl.int64)
+    eos = tl.load(query_start_loc_ptr + (pid_b + 1) * stride_qsl).to(tl.int64)
+    query_len = eos - bos
+    state_idx = tl.load(state_indices_ptr + pid_b * stride_si).to(tl.int64)
+    offs_k = tl.arange(0, BK)
+    offs_v = pid_v * BV + tl.arange(0, BV)
+    mask_k = offs_k < K
+    mask_v = offs_v < V
+    mask_state = mask_v[:, None] & mask_k[None, :]
+    if state_idx <= null_block_id:
+        for t in tl.static_range(SPEC_QUERY_LEN):
+            tl.store(
+                out_ptr + (bos + t) * stride_out_token + pid_hv * V + offs_v,
+                tl.zeros([BV], dtype=tl.float32),
+                mask=(t < query_len) & mask_v,
+            )
+        return
+    state = tl.load(
+        state_ptr
+        + state_idx * stride_state_block
+        + pid_hv * stride_state_head
+        + offs_v[:, None] * stride_state_v
+        + offs_k[None, :] * stride_state_k,
+        mask=mask_state,
+        other=0.0,
+    ).to(tl.float32)
+    neg_A = -tl.exp(tl.load(A_log_ptr + pid_hv).to(tl.float32))
+    dt_bias = tl.load(dt_bias_ptr + pid_hv).to(tl.float32)
+    rec_base = (
+        replay_ptr + state_idx * stride_replay_block + pid_hv * stride_replay_head
+    )
+    for t in tl.static_range(SPEC_QUERY_LEN):
+        valid = t < query_len
+        tok = bos + t
+        q = tl.load(
+            q_ptr + tok * stride_q_token + pid_h * K + offs_k,
+            mask=valid & mask_k,
+            other=0.0,
+        ).to(tl.float32)
+        k = tl.load(
+            k_ptr + tok * stride_k_token + pid_h * K + offs_k,
+            mask=valid & mask_k,
+            other=0.0,
+        ).to(tl.float32)
+        v = tl.load(
+            v_ptr + tok * stride_v_token + pid_hv * V + offs_v,
+            mask=valid & mask_v,
+            other=0.0,
+        ).to(tl.float32)
+        x = (
+            tl.load(a_ptr + tok * stride_a_token + pid_hv, mask=valid, other=0.0).to(
+                tl.float32
+            )
+            + dt_bias
+        )
+        sp = tl.where(
+            softplus_beta * x <= softplus_threshold,
+            (1 / softplus_beta) * tl.log(1 + tl.exp(softplus_beta * x)),
+            x,
+        )
+        g = neg_A * sp
+        beta = tl.sigmoid(
+            tl.load(b_ptr + tok * stride_b_token + pid_hv, mask=valid, other=0.0).to(
+                tl.float32
+            )
+        )
+        if USE_QK_L2NORM:
+            q = q * tl.rsqrt(tl.sum(q * q) + 1e-6)
+            k = k * tl.rsqrt(tl.sum(k * k) + 1e-6)
+        q = q * scale
+        new_state = state * tl.exp(g)
+        corr = (v - tl.sum(new_state * k[None, :], 1)) * beta
+        new_state = new_state + corr[:, None] * k[None, :]
+        state = tl.where(valid, new_state, state)
+        o = tl.sum(state * q[None, :], 1)
+        tl.store(
+            out_ptr + tok * stride_out_token + pid_hv * V + offs_v,
+            o.to(out_ptr.dtype.element_ty),
+            mask=valid & mask_v,
+        )
+        rec = rec_base + t * stride_replay_pos
+        tl.store(rec + offs_v * stride_replay_dim, corr, mask=valid & mask_v)
+        if pid_v == 0:
+            tl.store(rec + (V + offs_k) * stride_replay_dim, k, mask=valid & mask_k)
+            tl.store(rec + (V + K) * stride_replay_dim, g, mask=valid)
+
+
+def gdn_recoverssm_verify(
+    A_log: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    dt_bias: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    checkpoint_state: torch.Tensor,
+    replay_cache: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    state_indices: torch.Tensor,
+    spec_query_len: int,
+    scale: float | None = None,
+    softplus_beta: float = 1.0,
+    softplus_threshold: float = 20.0,
+    use_qk_l2norm_in_kernel: bool = True,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Verify a GDN speculative window against its checkpoint without modifying the
+    checkpoint. q, k: [1, T, H, K]; v: [1, T, HV, V]; a, b: [T, HV] (or [1, T, HV]);
+    checkpoint_state: [blocks, HV, V, K]; replay_cache: [blocks, HV, spec_query_len, V +
+    K + 1] fp32. Returns out [1, T, HV, V]."""
+    _require(
+        q.ndim == 4 and q.shape[0] == 1, "q must have shape [1, tokens, heads, dim]"
+    )
+    _, total, H, K = q.shape
+    _require(k.shape == q.shape, "q and k shapes differ")
+    _require(
+        v.ndim == 4 and tuple(v.shape[:2]) == (1, total),
+        "v must have shape [1, tokens, value heads, dim]",
+    )
+    HV, V = v.shape[-2], v.shape[-1]
+    _require(
+        H > 0 and HV >= H and HV % H == 0,
+        "value heads must be a positive multiple of the q/k heads",
+    )
+    _require(
+        a.numel() == total * HV and b.numel() == total * HV,
+        "gate or beta shape is incompatible",
+    )
+    a2 = a.reshape(total, HV)
+    b2 = b.reshape(total, HV)
+    _require(
+        a2.stride(1) == 1 and b2.stride(1) == 1,
+        "gate and beta heads must be contiguous",
+    )
+    _require(
+        tuple(A_log.shape) == (HV,) and tuple(dt_bias.shape) == (HV,),
+        "A_log or dt_bias shape is incompatible",
+    )
+    _require(
+        all(t_.stride(-1) == 1 and t_.stride(-2) == K for t_ in (q, k)),
+        "q and k heads must be contiguous",
+    )
+    _require(v.stride(-1) == 1 and v.stride(-2) == V, "v heads must be contiguous")
+    _require(checkpoint_state.ndim == 4, "checkpoint must be four-dimensional")
+    nb = checkpoint_state.shape[0]
+    _require(
+        tuple(checkpoint_state.shape[1:]) == (HV, V, K),
+        "checkpoint shape is incompatible",
+    )
+    _rec = (nb, HV, spec_query_len, V + K + 1)
+    _require(tuple(replay_cache.shape) == _rec, f"replay buffer needs shape {_rec}")
+    _require(replay_cache.dtype == torch.float32, "replay buffer must use float32")
+    _require(state_indices.ndim == 1, "state indices must be one-dimensional")
+    _require(
+        query_start_loc.ndim == 1
+        and query_start_loc.shape[0] == state_indices.shape[0] + 1,
+        "query metadata is incompatible",
+    )
+    _require(
+        total <= state_indices.shape[0] * spec_query_len,
+        "speculative decode input exceeds its activation capacity",
+    )
+    _dev = q.device
+    _require(
+        all(
+            t_.device == _dev
+            for t_ in (
+                k,
+                v,
+                a,
+                b,
+                A_log,
+                dt_bias,
+                checkpoint_state,
+                replay_cache,
+                query_start_loc,
+                state_indices,
+            )
+            if t_ is not None
+        )
+        and (out is None or out.device == _dev),
+        "inputs must be on the same device",
+    )
+    if scale is None:
+        scale = K**-0.5
+    if out is None:
+        out = torch.empty(1, total, HV, V, dtype=v.dtype, device=v.device)
+    _require(
+        tuple(out.shape) == (1, total, HV, V) and out.stride()[2:] == (V, 1),
+        "output shape or layout is incompatible",
+    )
+    batch = state_indices.shape[0]
+    if total == 0 or batch == 0:
+        return out
+    BK = triton.next_power_of_2(K)
+    BV = min(triton.next_power_of_2(V), 32)
+    grid = (triton.cdiv(V, BV), batch, HV)
+    _gdn_recoverssm_verify_kernel[grid](
+        q,
+        k,
+        v,
+        a2,
+        b2,
+        A_log,
+        dt_bias,
+        checkpoint_state,
+        replay_cache,
+        out,
+        query_start_loc,
+        state_indices,
+        scale,
+        softplus_beta,
+        softplus_threshold,
+        NULL_BLOCK_ID,
+        q.stride(1),
+        k.stride(1),
+        v.stride(1),
+        a2.stride(0),
+        b2.stride(0),
+        checkpoint_state.stride(0),
+        checkpoint_state.stride(1),
+        checkpoint_state.stride(2),
+        checkpoint_state.stride(3),
+        replay_cache.stride(0),
+        replay_cache.stride(1),
+        replay_cache.stride(2),
+        replay_cache.stride(3),
+        out.stride(1),
+        query_start_loc.stride(0),
+        state_indices.stride(0),
+        H=H,
+        HV=HV,
+        K=K,
+        V=V,
+        BK=BK,
+        BV=BV,
+        SPEC_QUERY_LEN=spec_query_len,
+        USE_QK_L2NORM=use_qk_l2norm_in_kernel,
+        num_warps=4,
+        num_stages=2,
+    )
+    return out
+
+
+@triton.jit
+def _commit_gdn_state_kernel(
+    state_ref_ptr,
+    state_base_addrs_ptr,
+    state_block_strides_ptr,
+    replay_ref_ptr,
+    replay_base_addrs_ptr,
+    replay_block_strides_ptr,
+    state_indices_ptr,
+    commit_lens_ptr,
+    final_state_indices_ptr,
+    boundary_state_indices_ptr,
+    boundary_recovery_lens_ptr,
+    null_block_id,
+    stride_state_head,
+    stride_state_v,
+    stride_state_k,
+    stride_replay_head,
+    stride_replay_pos,
+    stride_replay_dim,
+    stride_si,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    NUM_HEADS: tl.constexpr,
+    ALIGN_MODE: tl.constexpr,
+):
+    pid_v = tl.program_id(0)
+    pid_b = tl.program_id(1)
+    pid_lh = tl.program_id(2)
+    pid_l = pid_lh // NUM_HEADS
+    pid_h = pid_lh % NUM_HEADS
+    src = tl.load(state_indices_ptr + pid_b * stride_si).to(tl.int64)
+    if src <= null_block_id:
+        return
+    n = tl.load(commit_lens_ptr + pid_b)
+    if n == 0:
+        return
+    final_idx = tl.load(final_state_indices_ptr + pid_b).to(tl.int64)
+    boundary_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64)
+    boundary_len = tl.load(boundary_recovery_lens_ptr + pid_b)
+    if final_idx <= null_block_id:
+        return
+    state_ptr = tl.load(state_base_addrs_ptr + pid_l).to(
+        tl.pointer_type(state_ref_ptr.dtype.element_ty)
+    )
+    sbs = tl.load(state_block_strides_ptr + pid_l)
+    replay_ptr = tl.load(replay_base_addrs_ptr + pid_l).to(
+        tl.pointer_type(replay_ref_ptr.dtype.element_ty)
+    )
+    rbs = tl.load(replay_block_strides_ptr + pid_l)
+    rec_base = replay_ptr + src * rbs + pid_h * stride_replay_head
+    offs_k = tl.arange(0, BK)
+    offs_v = pid_v * BV + tl.arange(0, BV)
+    mask_k = offs_k < K
+    mask_v = offs_v < V
+    mask_state = mask_v[:, None] & mask_k[None, :]
+    sp = offs_v[:, None] * stride_state_v + offs_k[None, :] * stride_state_k
+    s0 = tl.load(
+        state_ptr + src * sbs + pid_h * stride_state_head + sp,
+        mask=mask_state,
+        other=0.0,
+    ).to(tl.float32)
+    fin_decay = 1.0
+    fin_corr = tl.zeros([BV, BK], tl.float32)
+    bnd_decay = 1.0
+    bnd_corr = tl.zeros([BV, BK], tl.float32)
+    for r in range(n):
+        t = n - r - 1
+        rec = rec_base + t * stride_replay_pos
+        c = tl.load(rec + offs_v * stride_replay_dim, mask=mask_v, other=0.0)
+        kk = tl.load(rec + (V + offs_k) * stride_replay_dim, mask=mask_k, other=0.0)
+        g = tl.load(rec + (V + K) * stride_replay_dim)
+        upd = c[:, None] * kk[None, :]
+        dec = tl.exp(g)
+        fin_corr += upd * fin_decay
+        fin_decay *= dec
+        if ALIGN_MODE:
+            before = t < boundary_len
+            bnd_corr += tl.where(before, upd * bnd_decay, 0.0)
+            bnd_decay *= tl.where(before, dec, 1.0)
+    if ALIGN_MODE:
+        bs = s0 * bnd_decay + bnd_corr
+        tl.store(
+            state_ptr + boundary_idx * sbs + pid_h * stride_state_head + sp,
+            bs.to(state_ref_ptr.dtype.element_ty),
+            mask=mask_state & (boundary_idx > null_block_id),
+        )
+    fs = s0 * fin_decay + fin_corr
+    tl.store(
+        state_ptr + final_idx * sbs + pid_h * stride_state_head + sp,
+        fs.to(state_ref_ptr.dtype.element_ty),
+        mask=mask_state,
+    )
+
+
+@dataclass
+class GDNRecoverSSMCommitContext:
+    conv_states: tuple[torch.Tensor, ...]
+    conv_state_base_addrs: torch.Tensor
+    conv_state_block_strides: torch.Tensor
+    conv_state_dim_strides: torch.Tensor
+    conv_state_token_strides: torch.Tensor
+    conv_history_len: int
+    checkpoints: tuple[torch.Tensor, ...]
+    state_base_addrs: torch.Tensor
+    state_block_strides: torch.Tensor
+    replays: tuple[torch.Tensor, ...]
+    replay_base_addrs: torch.Tensor
+    replay_block_strides: torch.Tensor
+    commit_lens: torch.Tensor
+    final_state_indices: torch.Tensor
+    boundary_state_indices: torch.Tensor
+    boundary_recovery_lens: torch.Tensor
+    spec_query_len: int
+
+    @classmethod
+    def from_tensors(
+        cls,
+        conv_states: Sequence[torch.Tensor],
+        checkpoints: Sequence[torch.Tensor],
+        replays: Sequence[torch.Tensor],
+        *,
+        spec_query_len: int,
+        max_num_reqs: int,
+        conv_dim_first: bool = True,
+    ) -> "GDNRecoverSSMCommitContext":
+        _require(len(checkpoints) > 0, "commit requires at least one layer")
+        _require(
+            len(conv_states) == len(checkpoints) == len(replays),
+            "conv, state and replay lists differ",
+        )
+        if not conv_dim_first:
+            conv_states = [s.transpose(-1, -2) for s in conv_states]
+        ref = checkpoints[0]
+        _require(ref.ndim == 4, "checkpoint must be four-dimensional")
+        nb, HV, V, K = ref.shape
+        _dev = ref.device
+        for s in checkpoints:
+            _require(
+                s.shape == ref.shape
+                and s.dtype == ref.dtype
+                and s.stride()[1:] == ref.stride()[1:]
+                and s.device == _dev,
+                "layers need matching checkpoints",
+            )
+        _rec = (nb, HV, spec_query_len, V + K + 1)
+        for r in replays:
+            _require(
+                tuple(r.shape) == _rec
+                and r.dtype == torch.float32
+                and r.stride()[1:] == replays[0].stride()[1:]
+                and r.device == _dev,
+                f"layers need matching float32 replay buffers of shape {_rec}",
+            )
+        conv_ref = conv_states[0]
+        _require(
+            conv_ref.ndim == 3 and conv_ref.shape[0] == nb,
+            "conv state must be [blocks, dim, window] with the checkpoint's "
+            "block count",
+        )
+        for c in conv_states:
+            _require(
+                c.shape == conv_ref.shape
+                and c.dtype == conv_ref.dtype
+                and c.stride()[1:] == conv_ref.stride()[1:]
+                and c.device == _dev,
+                "layers need matching conv states",
+            )
+        conv_dim, conv_len = conv_ref.shape[1:]
+        hist = conv_len - spec_query_len + 1
+        _require(hist > 0, "conv state is shorter than its window")
+        dev = ref.device
+        addr = lambda ts: torch.tensor(
+            [t.data_ptr() for t in ts], dtype=torch.int64, device=dev
+        )
+        bstr = lambda ts: torch.tensor(
+            [t.stride(0) for t in ts], dtype=torch.int64, device=dev
+        )
+        z = lambda: torch.empty(max_num_reqs, dtype=torch.int32, device=dev)
+        return cls(
+            tuple(conv_states),
+            addr(conv_states),
+            bstr(conv_states),
+            torch.tensor(
+                [s.stride(1) for s in conv_states], dtype=torch.int64, device=dev
+            ),
+            torch.tensor(
+                [s.stride(2) for s in conv_states], dtype=torch.int64, device=dev
+            ),
+            hist,
+            tuple(checkpoints),
+            addr(checkpoints),
+            bstr(checkpoints),
+            tuple(replays),
+            addr(replays),
+            bstr(replays),
+            z(),
+            z(),
+            z(),
+            z(),
+            spec_query_len,
+        )
+
+    @classmethod
+    def create(
+        cls, layers: Sequence[Any], *, spec_query_len: int, max_num_reqs: int
+    ) -> "GDNRecoverSSMCommitContext":
+        from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
+
+        if any(len(layer.kv_cache) != 3 for layer in layers):
+            raise ValueError("GDN RecoverSSM pages must contain conv, state and replay")
+        return cls.from_tensors(
+            [layer.kv_cache[0] for layer in layers],
+            [layer.kv_cache[1] for layer in layers],
+            [layer.kv_cache[2] for layer in layers],
+            spec_query_len=spec_query_len,
+            max_num_reqs=max_num_reqs,
+            conv_dim_first=is_conv_state_dim_first(),
+        )
+
+    def commit(
+        self,
+        num_accepted_tokens: torch.Tensor,
+        state_indices: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        request_indices: torch.Tensor | None = None,
+        block_table: torch.Tensor | None = None,
+        num_computed_tokens: torch.Tensor | None = None,
+        mamba_block_size: int | None = None,
+        commit_conv: bool = True,
+    ) -> None:
+        """Fold accepted GDN and convolution inputs into every layer's checkpoint."""
+        batch = state_indices.shape[0]
+        if batch == 0:
+            return
+        _require(state_indices.ndim == 1, "state indices must be one-dimensional")
+        _require(
+            batch <= self.commit_lens.shape[0], "commit batch exceeds its plan capacity"
+        )
+        _require(
+            query_start_loc.ndim == 1 and query_start_loc.shape[0] == batch + 1,
+            "commit metadata is incompatible",
+        )
+        _require(
+            request_indices is None or request_indices.shape[0] >= batch,
+            "request mapping is too short",
+        )
+        _require(
+            num_accepted_tokens.ndim == 1
+            and (request_indices is not None or num_accepted_tokens.shape[0] >= batch),
+            "accepted-token counts are too short",
+        )
+        _align = (block_table, num_computed_tokens, mamba_block_size)
+        _require(
+            all(x is None for x in _align) or all(x is not None for x in _align),
+            "align metadata is incomplete",
+        )
+        _require(
+            mamba_block_size is None or mamba_block_size >= self.spec_query_len,
+            "align block size must cover one speculative window",
+        )
+        _require(
+            block_table is None or block_table.ndim == 2,
+            "block table must be two-dimensional",
+        )
+        _dev = self.checkpoints[0].device
+        _require(
+            all(
+                t_ is None or t_.device == _dev
+                for t_ in (
+                    num_accepted_tokens,
+                    state_indices,
+                    query_start_loc,
+                    request_indices,
+                    block_table,
+                    num_computed_tokens,
+                )
+            ),
+            "commit inputs must be on the same device",
+        )
+        bt_stride = (0, 0) if block_table is None else block_table.stride()
+        prepare_commit_plan_kernel[(batch,)](
+            num_accepted_tokens,
+            request_indices,
+            state_indices,
+            query_start_loc,
+            block_table,
+            num_computed_tokens,
+            self.commit_lens,
+            self.final_state_indices,
+            self.boundary_state_indices,
+            self.boundary_recovery_lens,
+            NULL_BLOCK_ID,
+            mamba_block_size or 1,
+            block_table.shape[1] if block_table is not None else 1,
+            num_accepted_tokens.stride(0),
+            request_indices.stride(0) if request_indices is not None else 0,
+            state_indices.stride(0),
+            query_start_loc.stride(0),
+            bt_stride[0],
+            bt_stride[1],
+            0 if num_computed_tokens is None else num_computed_tokens.stride(0),
+            SPEC_QUERY_LEN=self.spec_query_len,
+            num_warps=1,
+        )
+        num_layers = len(self.checkpoints)
+        if commit_conv:
+            conv_ref = self.conv_states[0]
+            conv_dim = conv_ref.shape[1]
+            compact_conv_state_kernel[(triton.cdiv(conv_dim, 256), batch, num_layers)](
+                conv_ref,
+                self.conv_state_base_addrs,
+                self.conv_state_block_strides,
+                self.conv_state_dim_strides,
+                self.conv_state_token_strides,
+                state_indices,
+                self.commit_lens,
+                self.final_state_indices,
+                self.boundary_state_indices,
+                self.boundary_recovery_lens,
+                NULL_BLOCK_ID,
+                conv_dim,
+                self.conv_history_len,
+                state_indices.stride(0),
+                BLOCK_D=256,
+                BLOCK_HISTORY=triton.next_power_of_2(self.conv_history_len),
+                ALIGN_MODE=block_table is not None,
+                num_warps=4,
+            )
+        ref = self.checkpoints[0]
+        _, HV, V, K = ref.shape
+        BK = triton.next_power_of_2(K)
+        BV = min(triton.next_power_of_2(V), 32)
+        rr = self.replays[0]
+        _commit_gdn_state_kernel[(triton.cdiv(V, BV), batch, num_layers * HV)](
+            ref,
+            self.state_base_addrs,
+            self.state_block_strides,
+            rr,
+            self.replay_base_addrs,
+            self.replay_block_strides,
+            state_indices,
+            self.commit_lens,
+            self.final_state_indices,
+            self.boundary_state_indices,
+            self.boundary_recovery_lens,
+            NULL_BLOCK_ID,
+            ref.stride(1),
+            ref.stride(2),
+            ref.stride(3),
+            rr.stride(1),
+            rr.stride(2),
+            rr.stride(3),
+            state_indices.stride(0),
+            K=K,
+            V=V,
+            BK=BK,
+            BV=BV,
+            NUM_HEADS=HV,
+            ALIGN_MODE=block_table is not None,
+            num_warps=4,
+            num_stages=2,
+        )
+
+
+__all__ = ["GDNRecoverSSMCommitContext", "gdn_recoverssm_verify"]

vllm/model_executor/layers/mamba/ops/recoverssm.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Model-independent RecoverSSM Triton kernels shared by every recurrent-state
+implementation.
+
+RecoverSSM verifies a speculative window from a read-only checkpoint and, after
+sampling, commits the accepted prefix once. The per-model parts (the KDA and GDN verify
+kernels and their state reconstruction) live with their models; the kernels they share
+are here:
+
+- ``prepare_commit_plan_kernel``: per request, the commit length and the destination
+  state blocks (final and, in align mode, the block-boundary block) from the accepted
+  counts and the block table;
+- ``compact_conv_state_kernel``: moves the accepted suffix of an extended short-conv
+  window to the front of the destination block (and writes the boundary window in align
+  mode)."""
+
+from vllm.triton_utils import tl, triton
+
+
+@triton.heuristics(
+    {
+        "HAS_REQUEST_INDICES": lambda args: args["request_indices_ptr"] is not None,
+        "ALIGN_MODE": lambda args: args["block_table_ptr"] is not None,
+    }
+)
+@triton.jit
+def prepare_commit_plan_kernel(
+    num_accepted_ptr,
+    request_indices_ptr,
+    state_indices_ptr,
+    query_start_loc_ptr,
+    block_table_ptr,
+    num_computed_ptr,
+    commit_lens_ptr,
+    final_state_indices_ptr,
+    boundary_state_indices_ptr,
+    boundary_recovery_lens_ptr,
+    null_block_id,
+    mamba_block_size,
+    block_table_width,
+    stride_num_accepted,
+    stride_request_indices,
+    stride_state_indices,
+    stride_query_start_loc,
+    stride_block_table_row,
+    stride_block_table_col,
+    stride_num_computed,
+    SPEC_QUERY_LEN: tl.constexpr,
+    HAS_REQUEST_INDICES: tl.constexpr,
+    ALIGN_MODE: tl.constexpr,
+):
+    spec_idx = tl.program_id(0)
+    source_state_idx = tl.load(state_indices_ptr + spec_idx * stride_state_indices).to(
+        tl.int64
+    )
+    request_idx = spec_idx
+    if HAS_REQUEST_INDICES:
+        request_idx = tl.load(
+            request_indices_ptr + spec_idx * stride_request_indices
+        ).to(tl.int64)
+    num_accepted = tl.load(num_accepted_ptr + request_idx * stride_num_accepted).to(
+        tl.int32
+    )
+    bos = tl.load(query_start_loc_ptr + spec_idx * stride_query_start_loc).to(tl.int64)
+    eos = tl.load(query_start_loc_ptr + (spec_idx + 1) * stride_query_start_loc).to(
+        tl.int64
+    )
+    query_len = (eos - bos).to(tl.int32)
+    commit_len = tl.minimum(tl.maximum(num_accepted, 0), query_len)
+    commit_len = tl.minimum(commit_len, SPEC_QUERY_LEN)
+
+    final_state_idx = source_state_idx
+    boundary_state_idx = null_block_id
+    boundary_recovery_len = 0
+    if ALIGN_MODE:
+        num_computed = tl.load(num_computed_ptr + request_idx * stride_num_computed).to(
+            tl.int32
+        )
+        final_num_computed = num_computed + commit_len
+        # The committed state belongs to the block holding the last committed
+        # token. A window ending exactly on a boundary stays in that block; the
+        # next block may still be unallocated.
+        final_state_col = tl.minimum(
+            tl.maximum(final_num_computed - 1, 0) // mamba_block_size,
+            block_table_width - 1,
+        )
+        final_state_idx = tl.load(
+            block_table_ptr
+            + request_idx * stride_block_table_row
+            + final_state_col * stride_block_table_col
+        ).to(tl.int64)
+        next_boundary = (num_computed // mamba_block_size + 1) * mamba_block_size
+        crosses_boundary = final_num_computed >= next_boundary
+        boundary_recovery_len = next_boundary - num_computed
+        boundary_state_idx = tl.load(
+            block_table_ptr
+            + request_idx * stride_block_table_row
+            + (next_boundary // mamba_block_size - 1) * stride_block_table_col,
+            mask=crosses_boundary,
+            other=null_block_id,
+        ).to(tl.int64)
+    valid = (source_state_idx > null_block_id) & (commit_len > 0)
+    tl.store(commit_lens_ptr + spec_idx, tl.where(valid, commit_len, 0))
+    tl.store(
+        final_state_indices_ptr + spec_idx,
+        tl.where(valid, final_state_idx, null_block_id),
+    )
+    tl.store(
+        boundary_state_indices_ptr + spec_idx,
+        tl.where(valid, boundary_state_idx, null_block_id),
+    )
+    tl.store(
+        boundary_recovery_lens_ptr + spec_idx,
+        tl.where(valid, boundary_recovery_len, 0),
+    )
+
+
+@triton.jit
+def compact_conv_state_kernel(
+    conv_state_ref_ptr,
+    conv_state_base_addrs_ptr,
+    conv_state_block_strides_ptr,
+    conv_state_dim_strides_ptr,
+    conv_state_token_strides_ptr,
+    state_indices_ptr,
+    commit_lens_ptr,
+    final_state_indices_ptr,
+    boundary_state_indices_ptr,
+    boundary_recovery_lens_ptr,
+    null_block_id,
+    conv_dim,
+    conv_history_len,
+    stride_state_indices,
+    BLOCK_D: tl.constexpr,
+    BLOCK_HISTORY: tl.constexpr,
+    ALIGN_MODE: tl.constexpr,
+):
+    pid_d = tl.program_id(0)
+    pid_b = tl.program_id(1)
+    pid_l = tl.program_id(2)
+    source_state_idx = tl.load(state_indices_ptr + pid_b * stride_state_indices).to(
+        tl.int64
+    )
+    if source_state_idx <= null_block_id:
+        return
+
+    commit_len = tl.load(commit_lens_ptr + pid_b)
+    if commit_len == 0:
+        return
+    final_state_idx = tl.load(final_state_indices_ptr + pid_b).to(tl.int64)
+    boundary_state_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64)
+    boundary_recovery_len = tl.load(boundary_recovery_lens_ptr + pid_b)
+
+    if final_state_idx <= null_block_id:
+        return
+
+    base_addr = tl.load(conv_state_base_addrs_ptr + pid_l)
+    block_stride = tl.load(conv_state_block_strides_ptr + pid_l)
+    dim_stride = tl.load(conv_state_dim_strides_ptr + pid_l)
+    token_stride = tl.load(conv_state_token_strides_ptr + pid_l)
+    conv_state_ptr = base_addr.to(tl.pointer_type(conv_state_ref_ptr.dtype.element_ty))
+    source_ptr = conv_state_ptr + source_state_idx * block_stride
+    final_ptr = conv_state_ptr + final_state_idx * block_stride
+
+    offs_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D)
+    offs_h = tl.arange(0, BLOCK_HISTORY)
+    mask = (offs_d[:, None] < conv_dim) & (offs_h[None, :] < conv_history_len)
+    final_values = tl.load(
+        source_ptr
+        + offs_d[:, None] * dim_stride
+        + (commit_len - 1 + offs_h[None, :]) * token_stride,
+        mask=mask,
+    )
+    if ALIGN_MODE:
+        boundary_values = tl.load(
+            source_ptr
+            + offs_d[:, None] * dim_stride
+            + (boundary_recovery_len - 1 + offs_h[None, :]) * token_stride,
+            mask=mask & (boundary_state_idx > null_block_id),
+        )
+        boundary_ptr = conv_state_ptr + boundary_state_idx * block_stride
+        tl.store(
+            boundary_ptr
+            + offs_d[:, None] * dim_stride
+            + offs_h[None, :] * token_stride,
+            boundary_values,
+            mask=mask & (boundary_state_idx > null_block_id),
+        )
+    tl.store(
+        final_ptr + offs_d[:, None] * dim_stride + offs_h[None, :] * token_stride,
+        final_values,
+        mask=mask,
+    )
+
+
+__all__ = [
+    "compact_conv_state_kernel",
+    "prepare_commit_plan_kernel",
+]

vllm/model_executor/layers/mamba/recoverssm_utils.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Model-independent RecoverSSM Python helpers: the activation predicate shared by
+layers, models and backends, and argument validation that raises ``ValueError``
+before any raw-pointer kernel launch."""
+
+
+def recoverssm_require(cond: bool, msg: str, component: str = "RecoverSSM") -> None:
+    if not cond:
+        raise ValueError(f"{component}: {msg}")
+
+
+def uses_recoverssm(cache_config, num_speculative_tokens: int) -> bool:
+    """True when this model instance runs RecoverSSM for its recurrent states.
+
+    VllmConfig sets ``cache_config.use_recoverssm`` for every model whose recurrent
+    states support it (Kimi-K3 KDA, Qwen GDN/PLE). Without speculative tokens there is
+    nothing to recover, so the stock path runs."""
+    return (
+        bool(getattr(cache_config, "use_recoverssm", False))
+        and num_speculative_tokens > 0
+    )
+
+
+__all__ = ["recoverssm_require", "uses_recoverssm"]

vllm/model_executor/models/abliteration.py

@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Optional runtime refusal-direction removal ("abliteration").
+
+``VLLM_ABLITERATION`` (or the legacy ``QWEN38_ABLITERATION``) names a JSON file
+with a ``refusal_direction`` list, or a safetensors file with that tensor, of
+length ``hidden_size``. A bare name resolves to ``<name>.json`` in
+``/opt/qwen38/abliteration``, then in ``/data/abliteration``. When set, every
+write into the residual stream loses its component along the unit direction r:
+
+    x <- x - r (r . x)
+
+This runs on the attention block outputs (full and linear attention), the MLP
+output and the text token embeddings. For one direction it equals the weight
+edit ``W' = W - r (r^T W)`` on ``self_attn.o_proj``,
+``linear_attn.out_proj``, ``mlp.down_proj`` and the rows of ``embed_tokens``
+(Arditi et al., 2024), so the served weights, quantized or not, stay untouched.
+
+Unset, no projection kernel runs. The buffer always exists, holding zeros while
+the switch is off, so the attribute a compiled or captured graph refers to is
+always present, and the switch is part of the torch.compile cache key, so
+toggling it never reuses a stale graph.
+"""
+
+import json
+import math
+import os
+from functools import lru_cache
+
+import torch
+
+from vllm import envs
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+ENV = "VLLM_ABLITERATION"
+BUILTIN_DIRS = ("/opt/qwen38/abliteration", "/data/abliteration")
+DIRECTION = "_abliteration_r"
+ENABLED = "_abliteration_on"
+
+_shape_warned = False
+_register_logged = False
+
+
+def clear_logs() -> None:
+    """Forget the once-only log flags, for tests."""
+    global _register_logged, _shape_warned
+    _register_logged = False
+    _shape_warned = False
+
+
+def direction_path() -> str | None:
+    """The configured direction file, or None when the switch is off."""
+    value = (envs.VLLM_ABLITERATION or "").strip()
+    if value in ("", "0"):
+        return None
+    if os.sep in value or value.endswith((".json", ".safetensors")):
+        return value
+    for directory in BUILTIN_DIRS:
+        path = os.path.join(directory, f"{value}.json")
+        if os.path.exists(path):
+            return path
+    return os.path.join(BUILTIN_DIRS[0], f"{value}.json")
+
+
+@lru_cache(maxsize=4)
+def _read(path: str) -> torch.Tensor:
+    """The direction as a unit float32 CPU tensor, read once per path."""
+    if path.endswith(".json"):
+        with open(path) as f:
+            values = json.load(f)["refusal_direction"]
+        # A model built on the meta device would otherwise make this CPU tensor
+        # a meta one, and a meta tensor cannot be inspected or copied.
+        r = torch.tensor(values, dtype=torch.float32, device="cpu")
+    else:
+        from safetensors import safe_open
+
+        with safe_open(path, framework="pt", device="cpu") as f:
+            r = f.get_tensor("refusal_direction")
+        r = r.to(dtype=torch.float32, device="cpu")
+    r = r.flatten()
+    norm = float(r.norm())
+    if not math.isfinite(norm) or norm == 0.0:
+        raise ValueError(f"{ENV}: {path} holds a zero or non-finite direction")
+    logger.info("%s: refusal direction (%d values) from %s", ENV, r.numel(), path)
+    return r / norm
+
+
+def load_direction(hidden_size: int) -> torch.Tensor | None:
+    """The unit refusal direction on CPU, or None when the switch is off."""
+    path = direction_path()
+    if path is None:
+        return None
+    r = _read(path)
+    if r.numel() != hidden_size:
+        raise ValueError(
+            f"{ENV}: {path} holds {r.numel()} values; the model's hidden size "
+            f"is {hidden_size}"
+        )
+    return r
+
+
+@torch._dynamo.disable
+def register(module: torch.nn.Module, hidden_size: int, dtype: torch.dtype) -> None:
+    """Attach the direction as a non-persistent buffer, from a model constructor.
+
+    Call it before CUDA-graph capture so the buffer lands on the model's device.
+    It holds zeros when the switch is off. A construction pass that cannot hold
+    data, the meta device or a fake-tensor compile pass, turns the projection off
+    for that module and says so.
+    """
+    r = load_direction(hidden_size)
+    buf = torch.zeros(hidden_size, dtype=dtype)
+    on = r is not None
+    if on:
+        try:
+            buf.copy_(r)
+        except (RuntimeError, NotImplementedError):
+            # Meta or fake tensor construction: no room for a direction here.
+            logger.warning(
+                "%s: %s is constructed on the %s device, which cannot hold a "
+                "direction; projection is off for this module",
+                ENV,
+                type(module).__name__,
+                buf.device.type,
+            )
+            on = False
+    module.register_buffer(DIRECTION, buf, persistent=False)
+    setattr(module, ENABLED, on)
+    # Logging from project_ would sit inside a torch.compile region, which cannot
+    # trace a logger call, so the proof that the switch is live goes here.
+    global _register_logged
+    if on and not _register_logged:
+        _register_logged = True
+        logger.info(
+            "%s: projection is active, first registered on %s (hidden size %d)",
+            ENV,
+            type(module).__name__,
+            hidden_size,
+        )
+
+
+def project_(module: torch.nn.Module, x: torch.Tensor) -> torch.Tensor:
+    """Remove the component of x along the direction, in place. No-op when off."""
+    if not getattr(module, ENABLED, False):
+        return x
+    r = getattr(module, DIRECTION)
+    width = x.shape[-1]
+    if width != r.numel():
+        _warn_shape(type(module).__name__, width, r.numel())
+        return x
+    if r.dtype != x.dtype:
+        # Activations do not always carry the model dtype, and the rank-1 update
+        # needs one dtype throughout.
+        r = r.to(x.dtype)
+    if x.is_contiguous():
+        flat = x.view(-1, width)
+        flat.addr_(flat @ r, r, alpha=-1.0)
+    else:
+        x.sub_(torch.outer(x.reshape(-1, width) @ r, r).view(x.shape))
+    return x
+
+
+@torch._dynamo.disable
+def _warn_shape(name: str, width: int, size: int) -> None:
+    global _shape_warned
+    if not _shape_warned:
+        _shape_warned = True
+        logger.warning(
+            "%s: skipped a %d-wide tensor from %s; the direction has %d values",
+            ENV,
+            width,
+            name,
+            size,
+        )

vllm/model_executor/models/interfaces.py

@@ -88,6 +88,7 @@ MambaStateShapes: TypeAlias = (
     | tuple[tuple[int, int, int]]
     | tuple[tuple[int, int], tuple[int, int]]
     | tuple[tuple[int, int], tuple[int, int, int]]
+    | tuple[tuple[int, int], tuple[int, int, int], tuple[int, int, int]]
     | tuple[
         tuple[int, int],
         tuple[int, int, int],

vllm/model_executor/models/qwen3_5.py

@@ -61,11 +61,15 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
+from vllm.model_executor.model_loader.weight_utils import (
+    maybe_remap_kv_scale_name,
+)
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import MultiModalFeatureSpec
 from vllm.sequence import IntermediateTensors
 from vllm.tokenizers.registry import cached_tokenizer_from_config
 
+from . import abliteration
 from .interfaces import (
     HasInnerState,
     IsHybrid,
@@ -191,6 +195,7 @@ class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer):
         self.post_attention_layernorm = Qwen3_5RMSNorm(
             config.hidden_size, eps=config.rms_norm_eps
         )
+        abliteration.register(self, config.hidden_size, model_config.dtype)
 
         self.layer_scale = getattr(config, "layer_scale", False)
         if self.layer_scale:
@@ -251,7 +256,10 @@ class Qwen3_5Model(Qwen3NextModel):
         self.embed_tokens = VocabParallelEmbedding(
             self.vocab_size,
             config.hidden_size,
+            quant_config=self.quant_config,
+            prefix=f"{prefix}.embed_tokens",
         )
+        abliteration.register(self, config.hidden_size, vllm_config.model_config.dtype)
 
         def get_layer(prefix: str):
             return Qwen3_5DecoderLayer(
@@ -291,8 +299,28 @@ class Qwen3_5Model(Qwen3NextModel):
                 n_shared_experts=1,
                 ckpt_prefix="mlp.shared_expert",
             )
+        # Checkpoints that ship calibrated KV-cache scales name them
+        # `...self_attn.{k,v}_scale`, but the runtime parameters live on the
+        # inner attention layer (`...self_attn.attn.{k,v}_scale`), and
+        # AutoWeightsLoader performs no remapping — without this the scales
+        # are silently dropped and fp8 KV runs at scale 1.0 (#54623).
+        params_dict = dict(self.named_parameters())
+
+        def _remap_kv_scale_names(
+            weights: Iterable[tuple[str, torch.Tensor]],
+        ) -> Iterable[tuple[str, torch.Tensor]]:
+            for name, weight in weights:
+                if name.endswith((".k_scale", ".v_scale", ".q_scale")):
+                    remapped = maybe_remap_kv_scale_name(name, params_dict)
+                    if remapped is None:
+                        continue
+                    name = remapped
+                yield name, weight
+
         loader = AutoWeightsLoader(self)
-        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+        return loader.load_weights(
+            _remap_kv_scale_names(weights), mapper=self.hf_to_vllm_mapper
+        )
 
 
 class Qwen3_5ForCausalLMBase(

vllm/model_executor/models/qwen3_5_mtp.py

@@ -2,14 +2,16 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Inference-only Qwen3_5 MTP model."""
 
+import json
 from collections.abc import Iterable
 
+import huggingface_hub
 import torch
 from torch import nn
 from transformers import Qwen3_5MoeTextConfig, Qwen3_5TextConfig
 
 from vllm.compilation.decorators import support_torch_compile
-from vllm.config import VllmConfig
+from vllm.config import ModelConfig, VllmConfig
 from vllm.distributed import get_pp_group, tensor_model_parallel_all_gather
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.utils import (
@@ -33,7 +35,9 @@ from vllm.model_executor.models.qwen3_next import (
 )
 from vllm.model_executor.models.utils import sequence_parallel_chunk
 from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.repo_utils import try_get_local_file
 
+from . import abliteration
 from .interfaces import (
     MultiModalEmbeddings,
     SupportsMultiModal,
@@ -51,6 +55,64 @@ from .utils import (
 
 logger = init_logger(__name__)
 
+# Parameter suffixes a quantized MTP head would carry in the checkpoint.
+_MTP_QUANT_ARTIFACT_SUFFIXES = (
+    "weight_packed",
+    "weight_scale",
+    "weight_zero_point",
+    "input_scale",
+    "input_zero_point",
+)
+
+
+def _mtp_weight_names(model_config: ModelConfig) -> list[str] | None:
+    """Return the checkpoint's mtp.* tensor names, or None when unknown.
+
+    Reads the safetensors index from the local cache or the Hub, the same way
+    config-time helpers read recipe.yaml. A single-file checkpoint has no
+    index, so its layout stays unknown (callers keep the default behavior).
+    """
+    file_path = try_get_local_file(
+        model=model_config.model,
+        file_name="model.safetensors.index.json",
+        revision=model_config.revision,
+    )
+    if file_path is None:
+        try:
+            file_path = huggingface_hub.hf_hub_download(
+                model_config.model,
+                "model.safetensors.index.json",
+                revision=model_config.revision,
+            )
+        except Exception:
+            return None
+    try:
+        with open(file_path) as f:
+            weight_map = json.load(f).get("weight_map", {})
+    except (OSError, ValueError):
+        return None
+    return [name for name in weight_map if name.startswith("mtp.")]
+
+
+def _checkpoint_ships_dense_mtp(model_config: ModelConfig) -> bool:
+    """Whether the checkpoint ships the MTP head as dense weights.
+
+    Quantizers sometimes leave the MTP head unquantized without declaring it
+    in the quant config's ignore list; the tensors then arrive as plain
+    `mtp.*.weight` with no packed or scale companions (#58807). Detection is
+    name-based, so a properly quantized MTP head (any artifact suffix on any
+    mtp.* tensor) keeps the target's quantization.
+    """
+    names = _mtp_weight_names(model_config)
+    if not names:
+        return False
+    mtp_names = [name for name in names if name.startswith("mtp.")]
+    if not mtp_names:
+        return False
+    return not any(
+        name.rsplit(".", 1)[-1] in _MTP_QUANT_ARTIFACT_SUFFIXES for name in mtp_names
+    )
+
 
 @support_torch_compile(
     dynamic_arg_dims={
@@ -85,14 +147,25 @@ class Qwen3_5MultiTokenPredictor(nn.Module):
             self.vocab_size,
             config.hidden_size,
         )
+        abliteration.register(self, config.hidden_size, model_config.dtype)
 
         # Workaround: mtp.fc is stored as BF16 in NVFP4 checkpoints but is
         # missing from hf_quant_config.json exclude_modules. Force unquantized.
         # Ref: https://github.com/vllm-project/vllm/pull/38650
         # Ref: https://github.com/NVIDIA/Model-Optimizer/pull/1124
+        # compressed-tensors pack-quantized checkpoints may likewise ship the
+        # whole MTP head dense without declaring it in the ignore list (#58807).
+        mtp_dense = (
+            quant_config is not None
+            and quant_config.get_name() == "compressed-tensors"
+            and _checkpoint_ships_dense_mtp(model_config)
+        )
         fc_quant = (
             None
-            if (quant_config and quant_config.get_name() == "modelopt_fp4")
+            if (
+                quant_config
+                and (quant_config.get_name() == "modelopt_fp4" or mtp_dense)
+            )
             else quant_config
         )
         self.fc = ColumnParallelLinear(
@@ -108,13 +181,18 @@ class Qwen3_5MultiTokenPredictor(nn.Module):
         # GPTQ: quantized checkpoints may exclude MTP from quantization via
         # quantization_config.dynamic with "-:pattern" entries. When detected,
         # disable quantization for MTP layers so they use unquantized params.
+        # The dense-MTP detection above covers the same need for
+        # compressed-tensors checkpoints that declare nothing at all.
         original_quant = vllm_config.quant_config
         if quant_config and quant_config.get_name() not in ("modelopt_fp4",):
             hf_qc = getattr(model_config.hf_config, "quantization_config", None)
-            if isinstance(hf_qc, dict):
+            if mtp_dense or (
-                dynamic = hf_qc.get("dynamic", {})
+                isinstance(hf_qc, dict)
-                if any(k.startswith("-:") and "mtp" in k for k in dynamic):
+                and any(
-                    vllm_config.quant_config = None
+                    k.startswith("-:") and "mtp" in k for k in hf_qc.get("dynamic", {})
+                )
+            ):
+                vllm_config.quant_config = None
         self.layers = torch.nn.ModuleList(
             Qwen3_5DecoderLayer(
                 vllm_config,
@@ -141,7 +219,7 @@ class Qwen3_5MultiTokenPredictor(nn.Module):
         )
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
-        return self.embed_tokens(input_ids)
+        return abliteration.project_(self, self.embed_tokens(input_ids))
 
     def forward(
         self,

vllm/model_executor/models/qwen3_next.py

@@ -61,6 +61,7 @@ from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 from vllm.v1.attention.backend import AttentionType
 
+from . import abliteration
 from .interfaces import (
     EagleModelMixin,
     HasInnerState,
@@ -537,6 +538,7 @@ class Qwen3NextDecoderLayer(nn.Module):
         self.post_attention_layernorm = Qwen3NextRMSNorm(
             config.hidden_size, eps=config.rms_norm_eps
         )
+        abliteration.register(self, config.hidden_size, model_config.dtype)
 
         self.layer_scale = getattr(config, "layer_scale", False)
         if self.layer_scale:
@@ -584,6 +586,8 @@ class Qwen3NextDecoderLayer(nn.Module):
         else:
             raise ValueError("Invalid layer_type")
 
+        abliteration.project_(self, hidden_states)
+
         if self.layer_scale:
             if len(hidden_states.shape) == 2:
                 hidden_states = hidden_states * (
@@ -612,6 +616,8 @@ class Qwen3NextDecoderLayer(nn.Module):
         else:
             hidden_states = self.mlp(hidden_states)
 
+        abliteration.project_(self, hidden_states)
+
         if self.layer_scale:
             if len(hidden_states.shape) == 2:
                 hidden_states = hidden_states * (
@@ -661,6 +667,7 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):
             self.vocab_size,
             config.hidden_size,
         )
+        abliteration.register(self, config.hidden_size, vllm_config.model_config.dtype)
 
         def get_layer(prefix: str):
             return Qwen3NextDecoderLayer(
@@ -689,7 +696,7 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):
         self.aux_hidden_state_layers: tuple[int, ...] = ()
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
-        return self.embed_tokens(input_ids)
+        return abliteration.project_(self, self.embed_tokens(input_ids))
 
     @property
     def use_sequence_parallel(self) -> bool:

vllm/model_executor/models/qwen3_next_mtp.py

@@ -21,6 +21,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
+from vllm.model_executor.models import abliteration
 from vllm.model_executor.models.qwen3_next import (
     Qwen3NextDecoderLayer,
     Qwen3NextModel,
@@ -66,6 +67,7 @@ class Qwen3NextMultiTokenPredictor(nn.Module):
             self.vocab_size,
             config.hidden_size,
         )
+        abliteration.register(self, config.hidden_size, model_config.dtype)
 
         # Workaround: mtp.fc is stored as BF16 in NVFP4 checkpoints but is
         # missing from the checkpoint quant exclude list (its `ignore` glob
@@ -113,7 +115,7 @@ class Qwen3NextMultiTokenPredictor(nn.Module):
         )
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
-        return self.embed_tokens(input_ids)
+        return abliteration.project_(self, self.embed_tokens(input_ids))
 
     def forward(
         self,

vllm/model_executor/models/qwen3_vl.py

@@ -676,13 +676,25 @@ class Qwen3_VisionTransformer(nn.Module):
             ]
         )
 
+        # An offloader moves tower parameters to host RAM after this
+        # constructor returns, and no weight is left on the accelerator to
+        # read the compute device from. Record it here while they are still
+        # placed, so `device` keeps working once every weight is streamed.
+        first_param = next(self.parameters(), None)
+        self._compute_device = (
+            torch.device("cpu") if first_param is None else first_param.device
+        )
+
     @property
     def dtype(self) -> torch.dtype:
         return self.patch_embed.proj.weight.dtype
 
     @property
     def device(self) -> torch.device:
-        return self.patch_embed.proj.weight.device
+        weight_device = self.patch_embed.proj.weight.device
+        if weight_device.type != "cpu":
+            return weight_device
+        return self._compute_device
 
     def rot_pos_ids(self, h: int, w: int, spatial_merge_size: int) -> torch.Tensor:
         cache_key = (h, w, spatial_merge_size)
@@ -751,7 +763,7 @@ class Qwen3_VisionTransformer(nn.Module):
         for t, h, w in grid_thw:
             outputs.append(
                 interpolate_fn(
-                    self.pos_embed.weight,
+                    self.pos_embed_table(),
                     t,
                     h,
                     w,
@@ -762,6 +774,16 @@ class Qwen3_VisionTransformer(nn.Module):
             )
         return torch.cat(outputs, dim=0)
 
+    def pos_embed_table(self) -> torch.Tensor:
+        """The position embedding table on the compute device. The whole tower
+        may be streamed from host RAM, but the interpolation kernel cannot read
+        host memory.
+        """
+        table = self.pos_embed.weight
+        if table.device.type == "cpu":
+            table = table.to(self.device, non_blocking=True)
+        return table
+
     def prepare_encoder_metadata(
         self,
         grid_thw_list: Sequence[Sequence[int]],

vllm/model_executor/offloader/base.py

@@ -56,8 +56,9 @@ class BaseOffloader(ABC):
     `SupportsMultiModal._mark_tower_model`, outside the `make_layers` call.
 
     Offloaders whose `wrap_modules` may only be called on the decoder layer
-    stack (e.g. `PrefetchOffloader`, which schedules prefetches over a
+    stack must keep this `False`. `PrefetchOffloader` enables it per-instance
-    circular layer stack) must keep this `False`.
+    when `offload_towers` is configured, via a tower path that does not use
+    the circular layer stack or the static buffer pool.
     """
 
     @abstractmethod
@@ -143,8 +144,8 @@ def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader:
     """Create an offloader based on the offload configuration.
 
     Uses the explicit ``offload_backend`` selector.  When set to ``"auto"``,
-    selects prefetch if ``offload_group_size > 0``, UVA if
+    selects prefetch if ``offload_group_size > 0`` or ``offload_towers`` is
-    ``cpu_offload_gb > 0``, otherwise noop.
+    set, UVA if ``cpu_offload_gb > 0``, otherwise noop.
     """
     from vllm.model_executor.offloader.prefetch import PrefetchOffloader
     from vllm.model_executor.offloader.uva import UVAOffloader
@@ -154,7 +155,7 @@ def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader:
     prefetch = offload_config.prefetch
 
     if backend == "auto":
-        if prefetch.offload_group_size > 0:
+        if prefetch.offload_group_size > 0 or prefetch.offload_towers:
             backend = "prefetch"
         elif uva.cpu_offload_gb > 0:
             backend = "uva"
@@ -167,6 +168,7 @@ def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader:
             num_in_group=prefetch.offload_num_in_group,
             prefetch_step=prefetch.offload_prefetch_step,
             offload_params=prefetch.offload_params,
+            offload_towers=prefetch.offload_towers,
             mode="cpu",
         )
     elif backend == "uva":

vllm/model_executor/offloader/prefetch.py

@@ -150,6 +150,10 @@ class PrefetchOffloader(BaseOffloader):
         group_size: Group every N layers together.
         num_in_group: Offload this many layers per group (last N of each group).
         prefetch_step: Number of layers to prefetch ahead.
+        offload_towers: Also offload multimodal tower weights to pinned host
+            memory, streaming them to the GPU during each encoder forward
+            (bulk-copy). Group scheduling still applies to decoder layers
+            only; `group_size` may be 0 when only towers are offloaded.
         mode: Offload mode ("cpu" is currently supported).
 
     """
@@ -160,19 +164,24 @@ class PrefetchOffloader(BaseOffloader):
         num_in_group: int,
         prefetch_step: int,
         offload_params: set[str] | None = None,
+        offload_towers: bool = False,
         mode: str = "cpu",
     ):
         self.group_size = group_size
         self.num_in_group = num_in_group
         self.prefetch_step = prefetch_step
         self.offload_params = offload_params or set()
+        self.offload_towers = offload_towers
         self.mode = mode
+        if offload_towers:
+            self.supports_tower_offload = True
 
         # Copy stream for async H2D transfers
         self.copy_stream = torch.cuda.Stream()
 
         # Module offloaders and buffer pool (populated in wrap_modules/post_init)
         self.module_offloaders: list[_ModuleOffloader] = []
+        self.tower_streamers: list[_TowerUnitStreamer] = []
         self.buffer_pool: StaticBufferPool | None = None
         self.total_offloaded_bytes = 0
 
@@ -182,17 +191,29 @@ class PrefetchOffloader(BaseOffloader):
         prefix: str = "",
     ) -> list[nn.Module]:
         """Wrap modules with prefetch offloading logic."""
-        assert len(self.module_offloaders) == 0, (
+        modules = list(modules_generator)
-            "wrap_modules should only be called once"
-        )
 
         if prefix:
-            prefix = f"{prefix}."
+            # Routed by _mark_tower_model: a whole tower module, handled by
+            # the streaming path, independent of the decoder group schedule.
+            if self.offload_towers:
+                for module in modules:
+                    self._wrap_tower(module, prefix)
+            return modules
+
+        if self.group_size <= 0:
+            # Prefetch backend selected for towers only; decoder layers
+            # stay resident on the device.
+            return modules
+
+        assert len(self.module_offloaders) == 0, (
+            "wrap_modules should only be called once for the decoder stack"
+        )
 
         all_modules = []
         offload_modules = []
 
-        for module_index, module in enumerate(modules_generator):
+        for module_index, module in enumerate(modules):
             all_modules.append(module)
 
             # Select layers to offload based on group pattern
@@ -228,6 +249,68 @@ class PrefetchOffloader(BaseOffloader):
 
         return all_modules
 
+    def _wrap_tower(self, tower: nn.Module, prefix: str):
+        """Split a tower into parameter-bearing units and place each one
+        in pinned host RAM until its forward executes (bulk-copy).
+
+        A tower restricts this to its heavy compute subtrees by declaring
+        `tower_offload_children`. Tower code reads a small anchor weight to find
+        its compute device, and passes embedding weights straight to a kernel,
+        so weights outside that list stay on the device.
+        """
+        first_param = next(tower.parameters(), None)
+        if first_param is None:
+            # Tower was staged out (e.g. --limit-mm-per-prompt zero for all
+            # its modalities) or carries no parameters.
+            return
+        device = first_param.device
+        if device.type == "cpu":
+            return
+
+        declared = getattr(tower, "tower_offload_children", None)
+        children = dict(tower.named_children())
+        if declared is None:
+            roots: list[nn.Module] = [tower]
+            kept: list[str] = []
+        else:
+            roots = [c for name, c in children.items() if name in declared]
+            kept = [name for name in children if name not in declared]
+
+        num_units = 0
+        for root in roots:
+            for submodule in root.modules():
+                param_names = [
+                    name for name, _ in submodule.named_parameters(recurse=False)
+                ]
+                if not param_names:
+                    continue
+                streamer = _TowerUnitStreamer(submodule, param_names, device)
+                streamer.offload_to_host_now()
+                self.tower_streamers.append(streamer)
+                num_units += 1
+
+        logger.info(
+            "[PrefetchOffloader] Queued tower '%s' for host streaming: "
+            "%d parameter units, %d children kept on device (%s) (weights "
+            "move to the GPU only during encoder forwards).",
+            prefix,
+            num_units,
+            len(kept),
+            ", ".join(kept) or "none",
+        )
+
+        if isinstance(getattr(type(tower), "device", None), property):
+            resolved = getattr(tower, "device", None)
+            if isinstance(resolved, torch.device) and resolved.type != device.type:
+                logger.warning(
+                    "[PrefetchOffloader] Tower '%s' resolves its compute "
+                    "device to %s from offloaded weights, so code that calls "
+                    ".to(tower.device) runs on the host. Add the anchor "
+                    "submodule to tower_offload_children.",
+                    prefix,
+                    resolved,
+                )
+
     def _hook_module_forward(self, index: int, module: nn.Module):
         """Hook module's forward with torch.compile-compatible sync."""
         original_forward = module.forward
@@ -340,6 +423,21 @@ class PrefetchOffloader(BaseOffloader):
         (in _CpuParamOffloader.__init__), so GPU memory is available for the
         static buffer pool.
         """
+        # Tower streamers: re-snapshot host storage from the final processed
+        # weights, then install the streaming forwards.
+        for streamer in self.tower_streamers:
+            streamer.sync_storage()
+        for streamer in self.tower_streamers:
+            streamer.hook_forward()
+        if self.tower_streamers:
+            tower_bytes = sum(s.offloaded_bytes for s in self.tower_streamers)
+            self.total_offloaded_bytes += tower_bytes
+            logger.info_once(
+                f"[PrefetchOffloader] Tower offload active: "
+                f"{len(self.tower_streamers)} units, "
+                f"{tower_bytes / 1e9:.4f} GB in pinned host RAM."
+            )
+
         # Sync CPU storage with current param.data BEFORE collecting param info.
         # This is needed because process_weights_after_loading may have:
         # 1. Transformed weights (quantization, transpose, etc.)
@@ -391,6 +489,136 @@ class PrefetchOffloader(BaseOffloader):
             self.module_offloaders[i].start_onload_to_static()
 
 
+class _TowerUnitStreamer:
+    """Bulk-copy placement for one parameter-bearing submodule of a tower.
+
+    The unit's parameters live in pinned host memory between forwards.
+    On each forward, they are bulk-copied to freshly allocated GPU tensors
+    on the compute stream (so copies are stream-ordered ahead of the unit's
+    kernels with no explicit sync) and released right after. At steady
+    state only the unit currently executing holds VRAM, and freed blocks
+    are recycled by the caching allocator for the next unit's weights.
+
+    This is the bulk-copy alternative to UVA zero-copy for multimodal
+    towers: whole-tensor H2D transfers instead of operand tiles re-read
+    over PCIe inside the GEMM inner loop.
+
+    Interaction notes:
+    - If an encoder CUDA graph captures a streamed unit, correctness is
+      preserved (the memcpy is captured and replays from host memory), but
+      the captured allocations persist in the graph pool, so the VRAM
+      saving is lost. Prefer eager encoder execution (`cudagraph_mm_encoder`
+      off) together with tower offloading.
+    - Sleep/wake weight reload is not intercepted; storage is re-synced
+      from `param.data` at `post_init` only.
+    """
+
+    def __init__(self, unit: nn.Module, param_names: list[str], device: Any):
+        self.unit = unit
+        self.param_names = param_names
+        self.device = device
+        self.offloaded_bytes = 0
+        # name -> pinned host tensor holding the weights between forwards.
+        # Populated/refreshed in sync_storage() (after weight loading and
+        # process_weights_after_loading), which is the authoritative point.
+        self._storages: dict[str, torch.Tensor] = {}
+
+    def _get_param(self, name: str) -> nn.Parameter:
+        return getattr(self.unit, name)
+
+    @staticmethod
+    def _empty_cpu_like(data: torch.Tensor, pin_memory: bool) -> torch.Tensor:
+        return torch.empty_strided(
+            size=data.size(),
+            stride=data.stride(),
+            dtype=data.dtype,
+            layout=data.layout,
+            device="cpu",
+            pin_memory=pin_memory,
+        )
+
+    def offload_to_host_now(self) -> None:
+        """Move current parameter data to pinned host storage during
+        construction, so weight loading writes to host memory and the
+        device never accumulates tower weights.
+        """
+        pin_memory = should_pin_memory()
+        for name in self.param_names:
+            p = self._get_param(name)
+            if p.device.type == "cpu":
+                continue
+            cpu_data = self._empty_cpu_like(p.data, pin_memory)
+            cpu_data.copy_(p.data)
+            p.data = cpu_data
+
+    def sync_storage(self) -> None:
+        """Re-snapshot host storage from param.data.
+
+        Called from post_init() after weight loading and
+        process_weights_after_loading. device_loading_context may have
+        replaced param.data with fresh (already pinned) CPU tensors, and
+        quantization processing may have left data on the device; handle
+        both so _storages always references the final processed weights.
+        """
+        self.offloaded_bytes = 0
+        for name in list(self.param_names):
+            try:
+                p = self._get_param(name)
+            except AttributeError:
+                # Parameter deleted by process_weights_after_loading.
+                self.param_names.remove(name)
+                continue
+            data = p.data
+            if data.device.type != "cpu":
+                cpu_data = self._empty_cpu_like(data, pin_memory=should_pin_memory())
+                cpu_data.copy_(data)
+                data = cpu_data
+                p.data = data
+            elif should_pin_memory() and not data.is_pinned():
+                pinned = self._empty_cpu_like(data, pin_memory=True)
+                pinned.copy_(data)
+                data = pinned
+                p.data = data
+            self._storages[name] = data
+            self.offloaded_bytes += data.numel() * data.element_size()
+
+    def hook_forward(self) -> None:
+        if getattr(self.unit, "_vllm_tower_streamed", False):
+            return  # post_init ran more than once
+        original_forward = self.unit.forward
+        streamer = self
+
+        def forward(*args, **kwargs):
+            if torch.cuda.is_current_stream_capturing():
+                logger.warning_once(
+                    "[_TowerUnitStreamer] Streaming a tower unit during "
+                    "CUDA graph capture: results stay correct but the "
+                    "weights will be captured into the graph pool, so no "
+                    "VRAM is saved. Disable encoder CUDA graphs "
+                    "(cudagraph_mm_encoder) to keep the saving."
+                )
+            for name in streamer.param_names:
+                storage = streamer._storages.get(name)
+                if storage is None:
+                    continue
+                # Async H2D on the compute stream: ordered ahead of this
+                # unit's kernels, no explicit synchronization needed.
+                streamer._get_param(name).data = storage.to(
+                    streamer.device, non_blocking=True
+                )
+            try:
+                return original_forward(*args, **kwargs)
+            finally:
+                for name in streamer.param_names:
+                    storage = streamer._storages.get(name)
+                    if storage is None:
+                        continue
+                    streamer._get_param(name).data = storage
+
+        self.unit._vllm_tower_streamed = True
+        self.unit.forward = forward
+
+
 class _ModuleOffloader:
     """Manages offloading for a single module.
 

vllm/model_executor/warmup/qwen_triton_warmup.py

@@ -22,6 +22,11 @@ _QWEN_MODEL_TYPES = frozenset(
         "qwen3_5_text",
         "qwen3_5_moe",
         "qwen3_5_moe_text",
+        # Qwen4Exp runs its GDN prefill and QSA on other backends, but the
+        # kernels this file warms are the Triton ones around them -- the gated
+        # RMS norm, the causal conv, the post-conv preparation, the decode
+        # state update -- and those stay whichever prefill kernel is chosen.
+        # Left out of this set they compile on the first prompt instead.
         "qwen4_exp",
         "qwen4_exp_text",
     }

vllm/models/kimi_k3/nvidia/kda.py

@@ -556,7 +556,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
             self.cache_config.mamba_cache_dtype,
             self.cache_config.mamba_ssm_cache_dtype,
         )
-        if self.cache_config.use_kda_recoverssm:
+        if self.cache_config.use_recoverssm:
             return MambaStateDtypeCalculator.append_kda_recoverssm_record(
                 base_dtypes, self.model_config.dtype
             )
@@ -572,7 +572,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
             conv_kernel_size=self.conv_size,
             num_spec=self.num_spec,
         )
-        if self.cache_config.use_kda_recoverssm:
+        if self.cache_config.use_recoverssm:
             return MambaStateShapeCalculator.append_kda_recoverssm_record(
                 base_shapes,
                 self.num_heads,
@@ -591,7 +591,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
         run_gemm_rs_ar: bool = False,
     ) -> None:
         super().__init__(config, vllm_config, prefix)
-        self.use_recoverssm = self.cache_config.use_kda_recoverssm
+        self.use_recoverssm = self.cache_config.use_recoverssm
         if self.cache_config.use_replayssm and not self.use_recoverssm:
             raise ValueError(
                 "Kimi-K3 supports --use-replayssm only with speculative decoding"

vllm/models/kimi_k3/nvidia/kda_metadata.py

@@ -319,7 +319,7 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
             isinstance(additional_config, dict)
             and additional_config.get("kda_prefill_backend") == "flashinfer"
         )
-        self.use_recoverssm = vllm_config.cache_config.use_kda_recoverssm
+        self.use_recoverssm = vllm_config.cache_config.use_recoverssm
         self.spec_state_slots = 1 if self.use_recoverssm else self.num_spec + 1
         self.recoverssm_num_accepted_tokens: torch.Tensor | None = None
         self.recoverssm_context: KDARecoverSSMCommitContext | None = None

vllm/models/kimi_k3/nvidia/model.py

@@ -1496,7 +1496,7 @@ class KimiLinearForCausalLM(
             vllm_config.cache_config.mamba_cache_dtype,
             vllm_config.cache_config.mamba_ssm_cache_dtype,
         )
-        if vllm_config.cache_config.use_kda_recoverssm:
+        if vllm_config.cache_config.use_recoverssm:
             dtypes = MambaStateDtypeCalculator.append_kda_recoverssm_record(
                 dtypes, vllm_config.model_config.dtype
             )
@@ -1521,7 +1521,7 @@ class KimiLinearForCausalLM(
             conv_kernel_size=hf_config.linear_attn_config["short_conv_kernel_size"],
             num_spec=num_spec,
         )
-        if vllm_config.cache_config.use_kda_recoverssm:
+        if vllm_config.cache_config.use_recoverssm:
             return MambaStateShapeCalculator.append_kda_recoverssm_record(
                 shapes,
                 hf_config.linear_attn_config["num_heads"],

vllm/models/kimi_k3/nvidia/ops/recoverssm.py

@@ -9,6 +9,12 @@ from typing import Any
 import torch
 
 from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
+from vllm.model_executor.layers.mamba.ops.recoverssm import (
+    compact_conv_state_kernel as _compact_conv_state_kernel,
+)
+from vllm.model_executor.layers.mamba.ops.recoverssm import (
+    prepare_commit_plan_kernel as _prepare_commit_plan_kernel,
+)
 from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 
@@ -223,179 +229,6 @@ def _kda_recoverssm_verify_kernel(
             )
 
 
-@triton.heuristics(
-    {
-        "HAS_REQUEST_INDICES": lambda args: args["request_indices_ptr"] is not None,
-        "ALIGN_MODE": lambda args: args["block_table_ptr"] is not None,
-    }
-)
-@triton.jit
-def _prepare_commit_plan_kernel(
-    num_accepted_ptr,
-    request_indices_ptr,
-    state_indices_ptr,
-    query_start_loc_ptr,
-    block_table_ptr,
-    num_computed_ptr,
-    commit_lens_ptr,
-    final_state_indices_ptr,
-    boundary_state_indices_ptr,
-    boundary_recovery_lens_ptr,
-    null_block_id,
-    mamba_block_size,
-    block_table_width,
-    stride_num_accepted,
-    stride_request_indices,
-    stride_state_indices,
-    stride_query_start_loc,
-    stride_block_table_row,
-    stride_block_table_col,
-    stride_num_computed,
-    SPEC_QUERY_LEN: tl.constexpr,
-    HAS_REQUEST_INDICES: tl.constexpr,
-    ALIGN_MODE: tl.constexpr,
-):
-    spec_idx = tl.program_id(0)
-    source_state_idx = tl.load(state_indices_ptr + spec_idx * stride_state_indices).to(
-        tl.int64
-    )
-    request_idx = spec_idx
-    if HAS_REQUEST_INDICES:
-        request_idx = tl.load(
-            request_indices_ptr + spec_idx * stride_request_indices
-        ).to(tl.int64)
-    num_accepted = tl.load(num_accepted_ptr + request_idx * stride_num_accepted).to(
-        tl.int32
-    )
-    bos = tl.load(query_start_loc_ptr + spec_idx * stride_query_start_loc).to(tl.int64)
-    eos = tl.load(query_start_loc_ptr + (spec_idx + 1) * stride_query_start_loc).to(
-        tl.int64
-    )
-    query_len = (eos - bos).to(tl.int32)
-    commit_len = tl.minimum(tl.maximum(num_accepted, 0), query_len)
-    commit_len = tl.minimum(commit_len, SPEC_QUERY_LEN)
-
-    final_state_idx = source_state_idx
-    boundary_state_idx = null_block_id
-    boundary_recovery_len = 0
-    if ALIGN_MODE:
-        num_computed = tl.load(num_computed_ptr + request_idx * stride_num_computed).to(
-            tl.int32
-        )
-        final_num_computed = num_computed + commit_len
-        final_state_col = tl.minimum(
-            tl.maximum(final_num_computed - 1, 0) // mamba_block_size,
-            block_table_width - 1,
-        )
-        final_state_idx = tl.load(
-            block_table_ptr
-            + request_idx * stride_block_table_row
-            + final_state_col * stride_block_table_col
-        ).to(tl.int64)
-        next_boundary = (num_computed // mamba_block_size + 1) * mamba_block_size
-        crosses_boundary = final_num_computed >= next_boundary
-        boundary_recovery_len = next_boundary - num_computed
-        boundary_state_idx = tl.load(
-            block_table_ptr
-            + request_idx * stride_block_table_row
-            + (next_boundary // mamba_block_size - 1) * stride_block_table_col,
-            mask=crosses_boundary,
-            other=null_block_id,
-        ).to(tl.int64)
-    valid = (source_state_idx > null_block_id) & (commit_len > 0)
-    tl.store(commit_lens_ptr + spec_idx, tl.where(valid, commit_len, 0))
-    tl.store(
-        final_state_indices_ptr + spec_idx,
-        tl.where(valid, final_state_idx, null_block_id),
-    )
-    tl.store(
-        boundary_state_indices_ptr + spec_idx,
-        tl.where(valid, boundary_state_idx, null_block_id),
-    )
-    tl.store(
-        boundary_recovery_lens_ptr + spec_idx,
-        tl.where(valid, boundary_recovery_len, 0),
-    )
-
-
-@triton.jit
-def _compact_conv_state_kernel(
-    conv_state_ref_ptr,
-    conv_state_base_addrs_ptr,
-    conv_state_block_strides_ptr,
-    conv_state_dim_strides_ptr,
-    conv_state_token_strides_ptr,
-    state_indices_ptr,
-    commit_lens_ptr,
-    final_state_indices_ptr,
-    boundary_state_indices_ptr,
-    boundary_recovery_lens_ptr,
-    null_block_id,
-    conv_dim,
-    conv_history_len,
-    stride_state_indices,
-    BLOCK_D: tl.constexpr,
-    BLOCK_HISTORY: tl.constexpr,
-    ALIGN_MODE: tl.constexpr,
-):
-    pid_d = tl.program_id(0)
-    pid_b = tl.program_id(1)
-    pid_l = tl.program_id(2)
-    source_state_idx = tl.load(state_indices_ptr + pid_b * stride_state_indices).to(
-        tl.int64
-    )
-    if source_state_idx <= null_block_id:
-        return
-
-    commit_len = tl.load(commit_lens_ptr + pid_b)
-    if commit_len == 0:
-        return
-    final_state_idx = tl.load(final_state_indices_ptr + pid_b).to(tl.int64)
-    boundary_state_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64)
-    boundary_recovery_len = tl.load(boundary_recovery_lens_ptr + pid_b)
-
-    if final_state_idx <= null_block_id:
-        return
-
-    base_addr = tl.load(conv_state_base_addrs_ptr + pid_l)
-    block_stride = tl.load(conv_state_block_strides_ptr + pid_l)
-    dim_stride = tl.load(conv_state_dim_strides_ptr + pid_l)
-    token_stride = tl.load(conv_state_token_strides_ptr + pid_l)
-    conv_state_ptr = base_addr.to(tl.pointer_type(conv_state_ref_ptr.dtype.element_ty))
-    source_ptr = conv_state_ptr + source_state_idx * block_stride
-    final_ptr = conv_state_ptr + final_state_idx * block_stride
-
-    offs_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D)
-    offs_h = tl.arange(0, BLOCK_HISTORY)
-    mask = (offs_d[:, None] < conv_dim) & (offs_h[None, :] < conv_history_len)
-    final_values = tl.load(
-        source_ptr
-        + offs_d[:, None] * dim_stride
-        + (commit_len - 1 + offs_h[None, :]) * token_stride,
-        mask=mask,
-    )
-    if ALIGN_MODE:
-        boundary_values = tl.load(
-            source_ptr
-            + offs_d[:, None] * dim_stride
-            + (boundary_recovery_len - 1 + offs_h[None, :]) * token_stride,
-            mask=mask & (boundary_state_idx > null_block_id),
-        )
-        boundary_ptr = conv_state_ptr + boundary_state_idx * block_stride
-        tl.store(
-            boundary_ptr
-            + offs_d[:, None] * dim_stride
-            + offs_h[None, :] * token_stride,
-            boundary_values,
-            mask=mask & (boundary_state_idx > null_block_id),
-        )
-    tl.store(
-        final_ptr + offs_d[:, None] * dim_stride + offs_h[None, :] * token_stride,
-        final_values,
-        mask=mask,
-    )
-
-
 @triton.jit
 def _commit_kda_state_kernel(
     state_ref_ptr,

vllm/models/qwen4_exp/amd/model.py

@@ -77,6 +77,7 @@ from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
 from vllm.v1.kv_cache_interface import MambaSpec
 
 from ..common.hyperconnection import is_hc_sequence_parallel_enabled
+from ..common.mtp import make_mtp_hidden_buffer
 from .hyperconnection import GatedResidual, HyperConnectionConfig
 from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm
 from .ops.sequence_parallel import sp_reduce_scatter, sp_shard
@@ -467,25 +468,11 @@ class Qwen4ExpModel(nn.Module):
         else:
             self.hyper_connection_mixer = None
 
-        spec_config = vllm_config.speculative_config
+        self._mtp_hidden_buffer = make_mtp_hidden_buffer(
-        # MTP HC multi-stream outputs: when speculative method=="mtp" and the
+            vllm_config,
-        # model uses HC with hc_count>1, retain the pre-final-mixer multi-stream
+            config,
-        # hidden state [T, hc_count*H] so the MTP drafter can feed a real
+            is_last_rank=get_pp_group().is_last_rank,
-        # multi-stream backbone hidden on its first step (scheme A). Derived
-        # purely from config (NOT node identity) so P/D nodes stay consistent.
-        needs_mtp_hidden = (
-            spec_config is not None
-            and getattr(spec_config, "method", None) == "mtp"
-            and get_pp_group().is_last_rank
         )
-        if needs_mtp_hidden:
-            self._mtp_hidden_buffer = torch.empty(
-                vllm_config.scheduler_config.max_num_batched_tokens,
-                config.hc_count * config.hidden_size,
-                dtype=vllm_config.model_config.dtype,
-            )
-        else:
-            self._mtp_hidden_buffer = None
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.embed_tokens(input_ids)

vllm/models/qwen4_exp/amd/mtp.py

@@ -34,7 +34,11 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding,
 )
 from vllm.model_executor.model_loader.utils import configure_quant_config
-from vllm.model_executor.models.interfaces import SupportsHCSP, SupportsPP
+from vllm.model_executor.models.interfaces import (
+    LocalArgmaxMixin,
+    SupportsHCSP,
+    SupportsPP,
+)
 from vllm.model_executor.models.qwen3_5 import Qwen3_5Model
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
@@ -395,7 +399,13 @@ class Qwen4ExpMultiTokenPredictor(nn.Module):
         "hidden_states": 0,
     }
 )
-class Qwen4ExpMTP(nn.Module, SupportsHCSP, SupportsPP, Qwen4ExpMixtureOfExperts):
+class Qwen4ExpMTP(
+    LocalArgmaxMixin,
+    nn.Module,
+    SupportsHCSP,
+    SupportsPP,
+    Qwen4ExpMixtureOfExperts,
+):
     packed_modules_mapping = {
         "qkv_proj": ["q_proj", "k_proj", "v_proj"],
         "gate_up_proj": ["gate_proj", "up_proj"],

vllm/models/qwen4_exp/common/mtp.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared MTP drafter state for Qwen4Exp."""
+
+import torch
+from transformers import Qwen4ExpTextConfig
+
+from vllm.config import VllmConfig
+
+
+def make_mtp_hidden_buffer(
+    vllm_config: VllmConfig,
+    config: Qwen4ExpTextConfig,
+    *,
+    is_last_rank: bool,
+) -> torch.Tensor | None:
+    """Buffer holding the pre-final-mixer multi-stream hidden state, or None.
+
+    When speculative ``method == "mtp"`` the drafter feeds a real multi-stream
+    backbone hidden on its first step, so the last stage retains
+    ``[T, hc_count * H]``. Whether it is needed is derived from the config
+    alone, not from node identity, so prefill and decode nodes agree.
+    """
+    spec_config = vllm_config.speculative_config
+    needs_mtp_hidden = (
+        spec_config is not None
+        and getattr(spec_config, "method", None) == "mtp"
+        and is_last_rank
+    )
+    if not needs_mtp_hidden:
+        return None
+    return torch.empty(
+        vllm_config.scheduler_config.max_num_batched_tokens,
+        config.hc_count * config.hidden_size,
+        dtype=vllm_config.model_config.dtype,
+        device=vllm_config.device_config.device,
+    )

vllm/models/qwen4_exp/common/ngram_embedding.py

@@ -59,7 +59,7 @@ logger = init_logger(__name__)
 
 
 class Qwen4ExpPLEEmbedding(PLEVocabParallelEmbedding, ABC):
-    """ETP-sharded PLE table shared by device and pinned-host backends."""
+    """PLE table base shared by the device, pinned-host, and file-gather backends."""
 
     supports_prefetch: ClassVar[bool] = False
 
@@ -171,9 +171,9 @@ class Qwen4ExpPLEEmbedding(PLEVocabParallelEmbedding, ABC):
 
 
 class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase):
-    """Quantization interface shared by resident and pinned PLE tables."""
+    """Quantization interface shared by device, pinned, and file-gather PLE tables."""
 
-    # PLE post-load processing only validates scales in their current storage.
+    # PLE post-load processing validates scales in place and binds host file shards.
     requires_device_loading: bool = False
 
     @staticmethod
@@ -282,6 +282,12 @@ class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase):
             BLOCK_D=triton.next_power_of_2(row_bytes),
         )
 
+    def process_weights_after_loading(self, layer: nn.Module) -> None:
+        """Bind host-file-gather tables once every checkpoint shard is loaded."""
+        bind = getattr(layer, "bind_file_shards", None)
+        if bind is not None:
+            bind()
+
     @abstractmethod
     def dequantize(
         self,
@@ -371,6 +377,7 @@ class Qwen4ExpPLEFp8EmbeddingMethod(Qwen4ExpPLEEmbeddingMethod):
         sentinel = torch.finfo(torch.float32).min
         if torch.any(layer.weight_scale == sentinel):
             raise ValueError("FP8 PLE checkpoint is missing its global scale")
+        super().process_weights_after_loading(layer)
 
     def dequantize(
         self,
@@ -454,6 +461,7 @@ class Qwen4ExpPLENvFp4EmbeddingMethod(Qwen4ExpPLEEmbeddingMethod):
         sentinel = torch.finfo(torch.float32).min
         if torch.any(layer.weight_scale_2 == sentinel):
             raise ValueError("NVFP4 PLE checkpoint is missing its global scale")
+        super().process_weights_after_loading(layer)
 
     def lookup_dtype(self, layer: nn.Module) -> torch.dtype:
         return layer.params_dtype

vllm/models/qwen4_exp/common/qsa_cache.py

@@ -401,12 +401,53 @@ def build_qsa_metadata_triton(
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     """Build QSA side-cache and optional pre-indexer work metadata."""
     num_tokens = common_attn_metadata.num_actual_tokens
-    num_mapped_tokens = int(common_attn_metadata.query_start_loc_cpu[-1])
     token_to_req = token_to_req_buffer[:num_tokens]
     logical_positions = logical_positions_buffer[:num_tokens]
     visible_blocks = visible_blocks_buffer[:num_tokens]
     slot_mapping = slot_mapping_buffer[:num_tokens]
-    num_reqs = common_attn_metadata.query_start_loc.shape[0] - 1
+    if num_tokens == 0 and k_work_metadata_buffer is None:
+        return token_to_req, logical_positions, visible_blocks, slot_mapping
+    _launch_qsa_metadata_kernel(
+        common_attn_metadata.query_start_loc,
+        common_attn_metadata.seq_lens,
+        common_attn_metadata.slot_mapping,
+        common_attn_metadata.block_table_tensor,
+        token_to_req,
+        logical_positions,
+        visible_blocks,
+        slot_mapping,
+        k_work_metadata_buffer,
+        num_mapped_tokens=int(common_attn_metadata.query_start_loc_cpu[-1]),
+        storage_block_size=storage_block_size,
+        compress_ratio=compress_ratio,
+        circular_buffer_size=circular_buffer_size,
+        request_capacity=request_capacity,
+    )
+    if circular_buffer_size == 0 and compress_ratio == 1:
+        slot_mapping = common_attn_metadata.slot_mapping[:num_tokens]
+    return token_to_req, logical_positions, visible_blocks, slot_mapping
+
+
+def _launch_qsa_metadata_kernel(
+    query_start_loc: torch.Tensor,
+    seq_lens: torch.Tensor,
+    common_slot_mapping: torch.Tensor,
+    block_table: torch.Tensor,
+    token_to_req: torch.Tensor,
+    logical_positions: torch.Tensor,
+    visible_blocks: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    k_work_metadata_buffer: torch.Tensor | None,
+    *,
+    num_mapped_tokens: int,
+    storage_block_size: int,
+    compress_ratio: int,
+    circular_buffer_size: int,
+    request_capacity: int | None,
+) -> None:
+    """Fill QSA metadata in place; capture-safe given fixed shapes and scalars."""
+    num_tokens = token_to_req.shape[0]
+    num_reqs = query_start_loc.shape[0] - 1
     assert num_reqs > 0
 
     if k_work_metadata_buffer is not None:
@@ -420,10 +461,6 @@ def build_qsa_metadata_triton(
         request_scan_size = 1
         max_num_work = 0
 
-    if num_tokens == 0 and k_work_metadata_buffer is None:
-        return token_to_req, logical_positions, visible_blocks, slot_mapping
-
-    block_table = common_attn_metadata.block_table_tensor
     num_search_steps = int(math.ceil(math.log2(num_reqs)))
     work_search_steps = int(math.ceil(math.log2(num_reqs)))
     # The same grid covers token tiles and, for the compressed cache, work tiles.
@@ -432,9 +469,9 @@ def build_qsa_metadata_triton(
         cdiv(max_num_work, 256) if k_work_metadata_buffer is not None else 0
     )
     _build_qsa_metadata_kernel[(max(num_token_blocks, num_work_blocks, 1),)](
-        common_attn_metadata.query_start_loc,
+        query_start_loc,
-        common_attn_metadata.seq_lens,
+        seq_lens,
-        common_attn_metadata.slot_mapping,
+        common_slot_mapping,
         block_table,
         token_to_req,
         logical_positions,
@@ -459,9 +496,6 @@ def build_qsa_metadata_triton(
         WORK_BLOCK_SIZE=256,
         num_warps=4,
     )
-    if circular_buffer_size == 0 and compress_ratio == 1:
-        slot_mapping = common_attn_metadata.slot_mapping[:num_tokens]
-    return token_to_req, logical_positions, visible_blocks, slot_mapping
 
 
 def _build_qsa_metadata_torch(
@@ -584,6 +618,8 @@ class QSAForwardMetadata(AttentionMetadata):
     logical_positions: torch.Tensor
     visible_blocks: torch.Tensor
     k_work_metadata: torch.Tensor
+    common_slot_mapping: torch.Tensor
+    num_mapped_tokens: int
     num_actual_tokens: int
     num_decodes: int
     num_decode_tokens: int
@@ -600,6 +636,7 @@ class QSAMetadataBuilder(AttentionMetadataBuilder[QSAForwardMetadata]):
     """Build QSA metadata from vLLM's cache-group-specific common metadata."""
 
     _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
+    supports_draft_decode_metadata_update = HAS_TRITON
 
     def __init__(
         self,
@@ -711,6 +748,8 @@ class QSAMetadataBuilder(AttentionMetadataBuilder[QSAForwardMetadata]):
             logical_positions=logical_positions,
             visible_blocks=visible_blocks,
             k_work_metadata=k_work_metadata,
+            common_slot_mapping=common_attn_metadata.slot_mapping,
+            num_mapped_tokens=int(common_attn_metadata.query_start_loc_cpu[-1]),
             num_actual_tokens=num_tokens,
             num_decodes=num_decodes,
             num_decode_tokens=num_decode_tokens,
@@ -723,6 +762,29 @@ class QSAMetadataBuilder(AttentionMetadataBuilder[QSAForwardMetadata]):
             compress_ratio=self.compress_ratio,
         )
 
+    def update_draft_decode_metadata(self, metadata: QSAForwardMetadata) -> None:
+        if metadata.num_actual_tokens == 0:
+            return
+        build_k_work = not self.is_circular_buffer and self.compress_ratio != 1
+        _launch_qsa_metadata_kernel(
+            metadata.query_start_loc,
+            metadata.seq_lens,
+            metadata.common_slot_mapping,
+            metadata.block_table,
+            metadata.token_to_req,
+            metadata.logical_positions,
+            metadata.visible_blocks,
+            metadata.slot_mapping,
+            metadata.k_work_metadata if build_k_work else None,
+            num_mapped_tokens=metadata.num_mapped_tokens,
+            storage_block_size=self.storage_block_size,
+            compress_ratio=self.compress_ratio,
+            circular_buffer_size=(
+                self.kv_cache_spec.block_size if self.is_circular_buffer else 0
+            ),
+            request_capacity=self.request_capacity if build_k_work else None,
+        )
+
 
 class QSAStateBackend(AttentionBackend):
     """Key-only dummy backend for out-of-band QSA side-cache operations."""

vllm/models/qwen4_exp/nvidia/model.py

@@ -36,12 +36,14 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
 from vllm.model_executor.models.interfaces import (
     HasInnerState,
     IsHybrid,
+    MambaStateShapes,
     MixtureOfExperts,
     MultiModalEmbeddings,
     SupportsHCSP,
     SupportsLoRA,
     SupportsMRoPE,
     SupportsPP,
+    SupportsReplaySSM,
     _require_is_multimodal,
 )
 from vllm.model_executor.models.qwen3_5 import (
@@ -83,6 +85,7 @@ from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
 from vllm.v1.kv_cache_interface import MambaSpec
 
 from ..common.hyperconnection import is_hc_sequence_parallel_enabled
+from ..common.mtp import make_mtp_hidden_buffer
 from .hyperconnection import GatedResidual, HyperConnectionConfig
 from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm
 from .ops.cute_dsl.hc_down_silu import request_hc_down_silu_warmup
@@ -467,25 +470,11 @@ class Qwen4ExpModel(nn.Module):
         else:
             self.hyper_connection_mixer = None
 
-        spec_config = vllm_config.speculative_config
+        self._mtp_hidden_buffer = make_mtp_hidden_buffer(
-        # MTP HC multi-stream outputs: when speculative method=="mtp" and the
+            vllm_config,
-        # model uses HC with hc_count>1, retain the pre-final-mixer multi-stream
+            config,
-        # hidden state [T, hc_count*H] so the MTP drafter can feed a real
+            is_last_rank=get_pp_group().is_last_rank,
-        # multi-stream backbone hidden on its first step (scheme A). Derived
-        # purely from config (NOT node identity) so P/D nodes stay consistent.
-        needs_mtp_hidden = (
-            spec_config is not None
-            and getattr(spec_config, "method", None) == "mtp"
-            and get_pp_group().is_last_rank
         )
-        if needs_mtp_hidden:
-            self._mtp_hidden_buffer = torch.empty(
-                vllm_config.scheduler_config.max_num_batched_tokens,
-                config.hc_count * config.hidden_size,
-                dtype=vllm_config.model_config.dtype,
-            )
-        else:
-            self._mtp_hidden_buffer = None
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.embed_tokens(input_ids)
@@ -692,6 +681,7 @@ class Qwen4ExpForCausalLM(
     SupportsPP,
     Qwen4ExpMixtureOfExperts,
     IsHybrid,
+    SupportsReplaySSM,
 ):
     packed_modules_mapping = {
         "qkv_proj": ["q_proj", "k_proj", "v_proj", "indexer.index_qk_proj"],
@@ -808,20 +798,38 @@ class Qwen4ExpForCausalLM(
             num_spec=num_spec,
         )
 
+    @classmethod
+    def _uses_gdn_recoverssm(cls, vllm_config: VllmConfig) -> bool:
+        """One predicate for the replay record's shape AND dtype (they must
+        agree for the MambaSpec)."""
+        num_spec = (
+            vllm_config.speculative_config.num_speculative_tokens
+            if vllm_config.speculative_config
+            else 0
+        )
+        from vllm.model_executor.layers.mamba.recoverssm_utils import uses_recoverssm
+
+        return uses_recoverssm(vllm_config.cache_config, num_spec or 0)
+
     @classmethod
     def get_gdn_mamba_state_dtype_from_config(
         cls, vllm_config: VllmConfig
-    ) -> tuple[torch.dtype, torch.dtype]:
+    ) -> tuple[torch.dtype, ...]:
-        return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
+        dtypes: tuple[torch.dtype, ...] = (
-            vllm_config.model_config.dtype,
+            MambaStateDtypeCalculator.gated_delta_net_state_dtype(
-            vllm_config.cache_config.mamba_cache_dtype,
+                vllm_config.model_config.dtype,
-            vllm_config.cache_config.mamba_ssm_cache_dtype,
+                vllm_config.cache_config.mamba_cache_dtype,
+                vllm_config.cache_config.mamba_ssm_cache_dtype,
+            )
         )
+        if cls._uses_gdn_recoverssm(vllm_config):  # replay record, fp32
+            dtypes = (*dtypes, torch.float32)
+        return dtypes
 
     @classmethod
     def get_gdn_mamba_state_shape_from_config(
         cls, vllm_config: VllmConfig
-    ) -> tuple[tuple[int, int], tuple[int, int]]:
+    ) -> MambaStateShapes:
         parallel_config = vllm_config.parallel_config
         hf_config = vllm_config.model_config.hf_text_config
         tp_size = parallel_config.tensor_parallel_size
@@ -830,7 +838,7 @@ class Qwen4ExpForCausalLM(
             if vllm_config.speculative_config
             else 0
         )
-        return MambaStateShapeCalculator.gated_delta_net_state_shape(
+        shapes = MambaStateShapeCalculator.gated_delta_net_state_shape(
             tp_size,
             hf_config.linear_num_key_heads,
             hf_config.linear_num_value_heads,
@@ -839,18 +847,30 @@ class Qwen4ExpForCausalLM(
             hf_config.linear_conv_kernel_dim,
             num_spec,
         )
+        if cls._uses_gdn_recoverssm(
+            vllm_config
+        ):  # replay record [HV, spec_query_len, V + K + 1]
+            shapes = (
+                *shapes,
+                (
+                    hf_config.linear_num_value_heads // tp_size,
+                    num_spec + 1,
+                    hf_config.linear_value_head_dim + hf_config.linear_key_head_dim + 1,
+                ),
+            )
+        return shapes
 
     @classmethod
     def get_mamba_state_dtype_from_config(
         cls,
         vllm_config: VllmConfig,
-    ) -> tuple[torch.dtype, torch.dtype]:
+    ) -> tuple[torch.dtype, ...]:
         return cls.get_gdn_mamba_state_dtype_from_config(vllm_config)
 
     @classmethod
     def get_mamba_state_shape_from_config(
         cls, vllm_config: VllmConfig
-    ) -> tuple[tuple[int, int], tuple[int, int]]:
+    ) -> MambaStateShapes:
         return cls.get_gdn_mamba_state_shape_from_config(vllm_config)
 
     @classmethod
@@ -939,6 +959,7 @@ class Qwen4ExpForConditionalGeneration(
     HasInnerState,
     SupportsHCSP,
     Qwen4ExpMixtureOfExperts,
+    SupportsReplaySSM,
 ):
     """Qwen3-VL vision tower backed by the Qwen4Exp language model."""
 
@@ -1131,14 +1152,15 @@ class Qwen4ExpForConditionalGeneration(
     def get_mamba_state_dtype_from_config(
         cls,
         vllm_config: VllmConfig,
-    ) -> tuple[torch.dtype, torch.dtype]:
+    ) -> tuple[torch.dtype, ...]:
         return Qwen4ExpForCausalLM.get_mamba_state_dtype_from_config(vllm_config)
 
     @classmethod
-    def get_mamba_state_shape_from_config(
+    def get_mamba_state_shape_from_config(  # type: ignore[override]
         cls,
         vllm_config: VllmConfig,
-    ) -> tuple[tuple[int, int], tuple[int, int]]:
+    ) -> MambaStateShapes:
+        # GDN RecoverSSM adds a third state (the replay record) to Qwen3.5's two.
         return Qwen4ExpForCausalLM.get_mamba_state_shape_from_config(vllm_config)
 
     @classmethod

vllm/models/qwen4_exp/nvidia/model_state.py

@@ -13,6 +13,12 @@ from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
 from vllm.v1.worker.gpu.states import RequestState
 
+from .ngram_embedding import (
+    Qwen4ExpNGramEmbedding,
+    Qwen4ExpPLEFileGatherEmbedding,
+    stage_checkpoint_rows,
+)
+
 
 class Qwen4ExpModelState(MambaHybridModelState):
     """Add rollback-safe PLE n-gram context to the model inputs."""
@@ -63,6 +69,12 @@ class Qwen4ExpModelState(MambaHybridModelState):
             dtype=torch.int32,
             device=self.device,
         )
+        self.file_gather_modules = [
+            module
+            for module in model.modules()
+            if isinstance(module, Qwen4ExpNGramEmbedding)
+            and isinstance(module.ngram_embedding, Qwen4ExpPLEFileGatherEmbedding)
+        ]
 
     def _prepare_ngram_context(
         self,
@@ -106,9 +118,17 @@ class Qwen4ExpModelState(MambaHybridModelState):
         query_start_loc[: num_reqs_padded + 1].copy_(input_batch.query_start_loc)
         # Represent unused capacity as trailing zero-length requests.
         query_start_loc[num_reqs_padded + 1 :].copy_(input_batch.query_start_loc[-1])
+        ngram_context = self._prepare_ngram_context(input_batch, req_states)
+        if self.file_gather_modules:
+            stage_checkpoint_rows(
+                self.file_gather_modules,
+                input_batch.input_ids,
+                query_start_loc,
+                ngram_context,
+            )
         model_inputs.update(
             query_start_loc=query_start_loc,
-            ngram_context=self._prepare_ngram_context(input_batch, req_states),
+            ngram_context=ngram_context,
         )
         return model_inputs
 

vllm/models/qwen4_exp/nvidia/mtp.py

@@ -35,7 +35,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding,
 )
 from vllm.model_executor.model_loader.utils import configure_quant_config
-from vllm.model_executor.models.interfaces import SupportsHCSP, SupportsPP
+from vllm.model_executor.models.interfaces import (
+    LocalArgmaxMixin,
+    SupportsHCSP,
+    SupportsPP,
+    SupportsReplaySSM,
+)
 from vllm.model_executor.models.qwen3_5 import Qwen3_5Model
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
@@ -397,7 +402,14 @@ class Qwen4ExpMultiTokenPredictor(nn.Module):
         return loader.load_weights(weights, mapper=mapper)
 
 
-class Qwen4ExpMTP(nn.Module, SupportsHCSP, SupportsPP, Qwen4ExpMixtureOfExperts):
+class Qwen4ExpMTP(
+    LocalArgmaxMixin,
+    nn.Module,
+    SupportsHCSP,
+    SupportsPP,
+    Qwen4ExpMixtureOfExperts,
+    SupportsReplaySSM,
+):
     packed_modules_mapping = {
         "qkv_proj": ["q_proj", "k_proj", "v_proj"],
         "gate_up_proj": ["gate_proj", "up_proj"],

vllm/models/qwen4_exp/nvidia/ngram_embedding.py

@@ -1,7 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Qwen4Exp n-gram embeddings with device and pinned-host storage."""
+"""Qwen4Exp n-gram embeddings with device, pinned-host and host-file-gather storage."""
 
+import mmap
+import os
 from collections.abc import Iterable
 
 import torch
@@ -44,6 +46,157 @@ __all__ = [
 ]
 
 
+def _maps_entry(ptr: int) -> tuple[str, int]:
+    """Return the path and file offset of ``ptr`` from ``/proc/self/maps``."""
+    with open("/proc/self/maps") as maps:
+        for line in maps:
+            fields = line.split(maxsplit=5)
+            low, high = (int(bound, 16) for bound in fields[0].split("-"))
+            if low <= ptr < high:
+                path = fields[5].rstrip("\n") if len(fields) > 5 else ""
+                return path, ptr - low + int(fields[2], 16)
+    return "", 0
+
+
+class Qwen4ExpPLEFileGatherEmbedding(Qwen4ExpPLEEmbedding):
+    """PLE rows gathered on the host from the loader's file-backed shard views."""
+
+    def __init__(
+        self,
+        num_embeddings: int,
+        embedding_dim: int,
+        *,
+        params_dtype: torch.dtype,
+        padding_size: int,
+        prefix: str,
+        embedding_method: Qwen4ExpPLEEmbeddingMethod,
+        num_ngram_heads: int = 1,
+        max_total_tokens: int = 0,
+        data_parallel_rank: int = 0,
+        reduce_results: bool = True,
+    ) -> None:
+        super().__init__(
+            num_embeddings,
+            embedding_dim,
+            params_dtype=params_dtype,
+            padding_size=padding_size,
+            prefix=prefix,
+            embedding_method=embedding_method,
+            data_parallel_rank=data_parallel_rank,
+            reduce_results=reduce_results,
+        )
+        self._dummy = get_current_vllm_config().load_config.load_format == "dummy"
+        shape = (max_total_tokens, num_ngram_heads, embedding_dim)
+        dtype, pin = self.weight.dtype, self.weight.device.type != "cpu"
+        self._staging = torch.zeros(shape, dtype=dtype, device=self.weight.device)
+        self._host_ids = torch.empty(
+            shape[:2], dtype=torch.int64, device="cpu", pin_memory=pin
+        )
+        # Step N+1 writes these rows only after its stream sync, so after step N's H2D.
+        self._host_rows = torch.empty(shape, dtype=dtype, device="cpu", pin_memory=pin)
+        self._shards: dict[int, torch.Tensor] | None = {}
+        self._shard_size = 0
+        self._bound = False
+        self._fds: dict[str, int] = {}
+        self._fds_t = self._bases = torch.empty(0, dtype=torch.int64)
+
+    def allocate_embedding_weight(
+        self,
+        num_embeddings: int,
+        embedding_dim: int,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Allocate no rows; they stay in the checkpoint files."""
+        return torch.empty(0, embedding_dim, dtype=dtype)
+
+    def accept_checkpoint_shard(
+        self,
+        shard_index: int,
+        loaded_weight: torch.Tensor,
+        shard_size: int,
+    ) -> None:
+        """Keep a checkpoint shard as a zero-copy byte view of its storage."""
+        if self._bound or self._shards is None:
+            self._shards = None
+            raise RuntimeError("PLE host file gather does not support weight reload")
+        if loaded_weight.dtype != self.weight.dtype:
+            raise ValueError(
+                f"PLE shard dtype {loaded_weight.dtype} must match the embedding "
+                f"dtype {self.weight.dtype}; host file gather serves rows as stored"
+            )
+        if shard_index in self._shards:
+            raise ValueError(f"Duplicate PLE embedding shard {shard_index}")
+        self._shards[shard_index] = loaded_weight.view(torch.uint8)
+        self._shard_size = shard_size
+
+    def bind_file_shards(self) -> None:
+        """Check that file-backed shards cover every row; dummy loads serve zeros."""
+        assert self._shards is not None
+        if not self._dummy:
+            rows = sum(shard.shape[0] for shard in self._shards.values())
+            if rows != self.org_vocab_size:
+                raise ValueError(
+                    f"PLE shards cover {rows} of {self.org_vocab_size} rows"
+                )
+            self._fds_t, self._bases = torch.zeros(2, len(self._shards)).long()
+            for index, shard in self._shards.items():
+                path, offset = _maps_entry(shard.data_ptr())
+                if not path or not os.path.isfile(path):
+                    raise ValueError(
+                        f"PLE embedding shard {index} is not a file-backed view "
+                        f"({path or 'anonymous'})"
+                    )
+                if path not in self._fds:
+                    self._fds[path] = os.open(path, os.O_RDONLY)
+                self._fds_t[index], self._bases[index] = self._fds[path], offset
+        self._bound = True
+
+    def stage_rows(self, num_tokens: int) -> None:
+        """Gather the rows for ``_host_ids[:num_tokens]`` and copy them to staging."""
+        shards = self._shards
+        if shards is None or not self._bound:
+            raise RuntimeError("PLE host file gather is unbound or was reloaded")
+        ids = self._host_ids[:num_tokens].flatten()
+        out = self._host_rows[:num_tokens].flatten(0, 1).view(torch.uint8)
+        if not shards:
+            out.zero_()
+        else:
+            ids, order = ids.sort()
+            shard = ids // self._shard_size
+            local = ids - shard * self._shard_size
+            counts = torch.bincount(shard, minlength=len(shards)).tolist()
+            page, row_bytes = mmap.PAGESIZE, out.shape[1]
+            start = self._bases[shard] + local * row_bytes
+            first, last = start // page, (start + row_bytes - 1) // page
+            run = torch.ones_like(first, dtype=torch.bool)
+            run[1:] = (first[1:] > last[:-1] + 1) | (shard[1:] != shard[:-1])
+            fds, lows = self._fds_t[shard[run]].tolist(), first[run].tolist()
+            # Queue every page read before the gather's first serial fault.
+            for fd, lo, hi in zip(fds, lows, last[run.roll(-1)].tolist()):
+                os.posix_fadvise(
+                    fd, lo * page, (hi - lo + 1) * page, os.POSIX_FADV_WILLNEED
+                )
+            groups = zip(order.split(counts), local.split(counts))
+            for index, (rows, rows_local) in enumerate(groups):
+                if rows.numel():
+                    out.index_copy_(0, rows, shards[index].index_select(0, rows_local))
+        self._staging[:num_tokens].copy_(
+            self._host_rows[:num_tokens], non_blocking=True
+        )
+
+    def finalize_prefetch(self, total_tokens: int) -> torch.Tensor:
+        """Return the rows staged for ``total_tokens`` of this step."""
+        return self._staging[:total_tokens].flatten(-2)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Return the rows staged for this step."""
+        return self.finalize_prefetch(hidden_states.shape[0])
+
+    def start_prefetch(self, ngram_ids: torch.Tensor) -> None:
+        """Rows are staged in ``prepare_inputs``, so prefetch is a no-op."""
+        return None
+
+
 class Qwen4ExpNGramEmbedding(nn.Module):
     _MASK64 = (1 << 64) - 1
     _SPLITMIX_GAMMA = 0x9E3779B97F4A7C15
@@ -214,14 +367,29 @@ class Qwen4ExpNGramEmbedding(nn.Module):
         if params_dtype is None:
             params_dtype = torch.get_default_dtype()
         engram_config = get_current_vllm_config().engram_config
+        host_file_gather = bool(
+            engram_config is not None and engram_config.host_file_gather
+        )
+        if host_file_gather and isinstance(
+            embedding_quant_method, Qwen4ExpPLENvFp4EmbeddingMethod
+        ):
+            raise ValueError(
+                "host_file_gather serves PLE rows as stored and does not decode "
+                "NVFP4 block scales"
+            )
         embedding_cls = (
-            Qwen4ExpPLEPinnedHostEmbedding
+            Qwen4ExpPLEFileGatherEmbedding
+            if host_file_gather
+            else Qwen4ExpPLEPinnedHostEmbedding
             if engram_config is not None and engram_config.cpu_offload
             else Qwen4ExpPLEDeviceEmbedding
         )
-        # A TP reduce-scatter covers all vocabulary owners only when ETP=TP.
+        # A TP reduce-scatter covers all vocabulary owners only when ETP=TP. File
+        # gather stages complete rows on every rank, so reducing would scale them
+        # by the group size.
         self.use_reduce_scatter = (
             self.use_sequence_parallel
+            and not host_file_gather
             and get_etp_group().world_size == get_tp_group().world_size
         )
         self.ngram_embedding = embedding_cls(
@@ -385,8 +553,11 @@ class Qwen4ExpNGramEmbedding(nn.Module):
     ) -> torch.Tensor:
         """Return full embeddings or a token shard after combining vocabulary owners."""
         embedding = self.ngram_embedding
-        if embedding.supports_prefetch:
+        if embedding.supports_prefetch or isinstance(
-            # Prefetch contains full tokens even when hidden states are sharded.
+            embedding, Qwen4ExpPLEFileGatherEmbedding
+        ):
+            # Prefetched and staged rows cover every token, even when hidden
+            # states are sharded.
             embeddings = embedding.finalize_prefetch(total_tokens=input_ids.shape[0])
         else:
             ngram_ids = self.compute_ngram_ids(
@@ -474,6 +645,12 @@ class Qwen4ExpNGramEmbedding(nn.Module):
                         f"{suffix}: expected {expected_shape}, got "
                         f"{tuple(loaded_weight.shape)}"
                     )
+                if isinstance(embedding, Qwen4ExpPLEFileGatherEmbedding):
+                    embedding.accept_checkpoint_shard(
+                        shard_index, loaded_weight, shard_size
+                    )
+                    loaded.add(f"ngram_embedding.{suffix}")
+                    continue
                 # A cast would silently reinterpret packed E2M1 bytes.
                 if nvfp4_method is not None and loaded_weight.dtype != parameter.dtype:
                     raise ValueError(
@@ -500,3 +677,20 @@ class Qwen4ExpNGramEmbedding(nn.Module):
         if regular_weights:
             loaded.update(AutoWeightsLoader(self).load_weights(regular_weights))
         return loaded
+
+
+def stage_checkpoint_rows(
+    modules: list[Qwen4ExpNGramEmbedding],
+    input_ids: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    ngram_context: torch.Tensor,
+) -> None:
+    """Stage host-file-gather PLE rows for one step behind a single stream sync."""
+    num_tokens = input_ids.shape[0]
+    for module in modules:
+        ids = module.compute_ngram_ids(input_ids, query_start_loc, ngram_context)
+        module.ngram_embedding._host_ids[:num_tokens].copy_(ids, non_blocking=True)
+    if input_ids.device.type != "cpu":
+        torch.accelerator.current_stream().synchronize()
+    for module in modules:
+        module.ngram_embedding.stage_rows(num_tokens)

vllm/models/qwen4_exp/nvidia/ple_layer.py

@@ -171,6 +171,14 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
         return True
 
     def get_attn_backend(self) -> type[PleShortConvAttentionBackend]:
+        from vllm.model_executor.layers.mamba.recoverssm_utils import uses_recoverssm
+
+        if uses_recoverssm(self.cache_config, self.num_spec_tokens):
+            from vllm.v1.attention.backends.ple_recoverssm import (
+                PleRecoverSSMAttentionBackend,
+            )
+
+            return PleRecoverSSMAttentionBackend
         return PleShortConvAttentionBackend
 
     def get_state_dtype(self) -> tuple[torch.dtype, ...]:

vllm/v1/attention/backends/gdn_recoverssm.py

@@ -0,0 +1,181 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GDN attention metadata for RecoverSSM speculative decode (Qwen-style Gated DeltaNet).
+
+Wraps the generic GDN builder: every active decode row runs through the spec path (a
+draft-less step is a T=1 window), the spec state indices keep only the checkpoint
+column, num_accepted_tokens is always 1 (the checkpoint and the compacted conv window
+already hold the accepted state), and the metadata carries what the post-sampling commit
+needs. Mirrors the Kimi-K3 KDA builder's RecoverSSM handling."""
+
+from dataclasses import dataclass, field, fields
+from typing import Any
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata
+from vllm.v1.attention.backends.gdn_attn import (
+    GDNAttentionBackend,
+    GDNAttentionMetadata,
+    GDNAttentionMetadataBuilder,
+)
+from vllm.v1.attention.backends.recoverssm_metadata import (
+    RecoverSSMMetadata,
+    RecoverSSMPostprocessMetadata,
+    recoverssm_request_indices,
+    recoverssm_spec_rows,
+)
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class GDNRecoverSSMCommitMetadata:
+    state_indices: torch.Tensor  # [num_spec_decodes, 1]
+    query_start_loc: torch.Tensor  # [num_spec_decodes + 1]
+    request_indices: (
+        torch.Tensor | None
+    )  # batch rows of the spec requests (None: rows 0..n-1)
+    block_table: torch.Tensor | None  # align mode
+    num_computed_tokens: torch.Tensor | None
+    block_size: int | None
+
+
+@dataclass
+class GDNRecoverSSMMetadata(GDNAttentionMetadata, RecoverSSMMetadata):
+    recoverssm_commit: GDNRecoverSSMCommitMetadata | None = None
+    recoverssm_context: Any = field(default=None, repr=False, compare=False)
+
+    def commit_recoverssm_state(
+        self, num_accepted_tokens: torch.Tensor
+    ) -> RecoverSSMPostprocessMetadata | None:
+        c = self.recoverssm_commit
+        if c is None:
+            return None
+        n = self.num_spec_decodes
+        self.recoverssm_context.commit(
+            num_accepted_tokens,
+            c.state_indices[:n, 0],
+            c.query_start_loc[: n + 1],
+            request_indices=c.request_indices,
+            block_table=c.block_table,
+            num_computed_tokens=c.num_computed_tokens,
+            mamba_block_size=c.block_size,
+        )
+        if c.block_table is None:
+            return None
+        assert c.block_size is not None and c.num_computed_tokens is not None
+        return RecoverSSMPostprocessMetadata(
+            num_spec_decodes=n,
+            request_indices=c.request_indices,
+            block_table=c.block_table,
+            num_computed_tokens=c.num_computed_tokens,
+            block_size=c.block_size,
+        )
+
+
+class GDNRecoverSSMMetadataBuilder(GDNAttentionMetadataBuilder):
+    # The verify/commit metadata is not staged for FULL graph replay yet; declaring
+    # NEVER makes the runner fall back to PIECEWISE (attention compiled piecewise).
+    _cudagraph_support = AttentionCGSupport.NEVER
+
+    def __init__(self, kv_cache_spec, layer_names, vllm_config, device) -> None:
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+        # The builder may be created before the runner resolves the CUDA graph mode.
+        self.use_full_cuda_graph = False
+        max_reqs = vllm_config.scheduler_config.max_num_seqs
+        self._ones = torch.ones(max_reqs, dtype=torch.int32, device=device)
+        self._context: Any = None
+        self._logged = False
+
+    def _get_context(self):
+        if self._context is None:
+            from vllm.model_executor.layers.mamba.gdn.recoverssm_gdn import (
+                GDNRecoverSSMCommitContext,
+            )
+
+            fc = self.vllm_config.compilation_config.static_forward_context
+            layers = [fc[name] for name in self.layer_names]
+            self._context = GDNRecoverSSMCommitContext.create(
+                layers,
+                spec_query_len=1 + self.num_spec,
+                max_num_reqs=self.vllm_config.scheduler_config.max_num_seqs,
+            )
+            logger.info(
+                "GDN RecoverSSM commit context: %d layers, spec_query_len %d",
+                len(layers),
+                1 + self.num_spec,
+            )
+        return self._context
+
+    def build(  # type: ignore[override]
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        num_accepted_tokens: torch.Tensor | None = None,
+        num_decode_draft_tokens_cpu: torch.Tensor | None = None,
+        fast_build: bool = False,
+    ) -> "GDNRecoverSSMMetadata":
+        m = common_attn_metadata
+        rows = None
+        if self.use_spec_decode and num_decode_draft_tokens_cpu is not None:
+            if m.is_prefilling is None or m.is_prefilling.device.type != "cpu":
+                raise ValueError("GDN RecoverSSM needs the CPU is_prefilling mask")
+            # Positive counts only defeat the base builder's "no drafts -> regular
+            # decode" shortcut: RecoverSSM keeps every post-prefill row on the spec path
+            # (extended conv window).
+            num_decode_draft_tokens_cpu, rows = recoverssm_spec_rows(
+                m.is_prefilling,
+                m.query_start_loc_cpu,
+                num_decode_draft_tokens_cpu,
+                m.num_reqs,
+            )
+        meta = super().build(
+            common_prefix_len,
+            m,
+            num_accepted_tokens=num_accepted_tokens,
+            num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu,
+            fast_build=fast_build,
+        )
+        base = {f.name: getattr(meta, f.name) for f in fields(meta)}
+        commit = None
+        if meta.num_spec_decodes > 0:
+            n = meta.num_spec_decodes
+            assert meta.spec_state_indices_tensor is not None
+            base["spec_state_indices_tensor"] = meta.spec_state_indices_tensor[:, :1]
+            base["num_accepted_tokens"] = self._ones[:n]
+            if rows is None:
+                raise ValueError("GDN RecoverSSM: spec rows without draft counts")
+            request_indices = recoverssm_request_indices(
+                rows, n, m.query_start_loc.device
+            )
+            align = self.vllm_config.cache_config.mamba_cache_mode == "align"
+            commit = GDNRecoverSSMCommitMetadata(
+                state_indices=base["spec_state_indices_tensor"],
+                query_start_loc=meta.spec_query_start_loc,
+                request_indices=request_indices,
+                block_table=m.block_table_tensor if align else None,
+                num_computed_tokens=m.compute_num_computed_tokens() if align else None,
+                block_size=self.kv_cache_spec.block_size if align else None,
+            )
+            if not self._logged:
+                self._logged = True
+                logger.info(
+                    "GDN RecoverSSM path taken: %d spec rows, align=%s", n, align
+                )
+        return GDNRecoverSSMMetadata(
+            **base,
+            recoverssm_commit=commit,
+            recoverssm_context=self._get_context() if commit is not None else None,
+        )
+
+
+class GDNRecoverSSMAttentionBackend(GDNAttentionBackend):
+    @staticmethod
+    def get_name() -> str:
+        return "GDN_RECOVERSSM"
+
+    @staticmethod
+    def get_builder_cls() -> type[GDNRecoverSSMMetadataBuilder]:
+        return GDNRecoverSSMMetadataBuilder

vllm/v1/attention/backends/mamba_attn.py

@@ -104,7 +104,12 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
         self.compilation_config = vllm_config.compilation_config
         self.num_spec_tokens: int = vllm_config.num_speculative_tokens
         self.use_spec_decode = self.num_spec_tokens > 0
-        self.use_replayssm = vllm_config.cache_config.use_replayssm
+        # RecoverSSM also rides on --use-replayssm but keeps its own metadata; the
+        # Mamba2 ReplaySSM ring-buffer workspaces below are not used then.
+        self.use_replayssm = (
+            vllm_config.cache_config.use_replayssm
+            and not vllm_config.cache_config.use_recoverssm
+        )
         self.replayssm_buffer_len = vllm_config.cache_config.replayssm_buffer_len
         self.use_flashinfer_replayssm = (
             self.use_replayssm

vllm/v1/attention/backends/ple_recoverssm.py

@@ -0,0 +1,326 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""RecoverSSM protocol for the Qwen4Exp PLE short-conv layer.
+
+When the GDN layers run RecoverSSM, the runner resets the shared num_accepted_tokens to
+1 after each commit (align mode). The PLE dilated short conv keeps an extended window in
+one state block and reads it at num_accepted - 1, so it must follow the same protocol:
+every active decode row takes the spec path, the conv reads at offset 0, and after
+sampling the accepted window is compacted to the front (the shared RecoverSSM compaction
+kernel; in align mode it also writes the block-boundary window)."""
+
+from dataclasses import dataclass, field, fields
+from typing import Any
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.model_executor.layers.mamba.ops.recoverssm import (
+    compact_conv_state_kernel,
+    prepare_commit_plan_kernel,
+)
+from vllm.model_executor.layers.mamba.recoverssm_utils import recoverssm_require
+from vllm.triton_utils import triton
+from vllm.v1.attention.backend import AttentionCGSupport
+from vllm.v1.attention.backends.recoverssm_metadata import (
+    RecoverSSMMetadata,
+    RecoverSSMPostprocessMetadata,
+    recoverssm_request_indices,
+    recoverssm_spec_rows,
+)
+from vllm.v1.attention.backends.short_conv_attn import (
+    PleShortConvAttentionBackend,
+    PleShortConvAttentionMetadata,
+    PleShortConvAttentionMetadataBuilder,
+)
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
+
+logger = init_logger(__name__)
+
+
+def _require(cond: bool, msg: str) -> None:
+    recoverssm_require(cond, msg, "PLE RecoverSSM")
+
+
+class _PleConvCommit:
+    def __init__(self, conv_states, spec_query_len: int, max_num_reqs: int):
+        from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
+
+        _require(len(conv_states) > 0, "PLE commit requires at least one layer")
+        if not is_conv_state_dim_first():
+            conv_states = [
+                s.transpose(-1, -2) for s in conv_states
+            ]  # -> [blocks, C, W]
+        self.conv_states = list(conv_states)
+        ref = self.conv_states[0]
+        _require(ref.ndim == 3, "PLE conv state must be [blocks, dim, window]")
+        dev = ref.device
+        for s in self.conv_states:
+            _require(
+                s.shape == ref.shape
+                and s.dtype == ref.dtype
+                and s.stride()[1:] == ref.stride()[1:]
+                and s.device == dev,
+                "PLE layers need matching conv states",
+            )
+        self.conv_dim, conv_len = ref.shape[1], ref.shape[2]
+        self.hist = conv_len - spec_query_len + 1
+        _require(self.hist > 0, "PLE conv state is shorter than its window")
+        self.spec_query_len = spec_query_len
+        t = lambda xs: torch.tensor(xs, dtype=torch.int64, device=dev)
+        self.base = t([s.data_ptr() for s in self.conv_states])
+        self.bstride = t([s.stride(0) for s in self.conv_states])
+        self.dstride = t([s.stride(1) for s in self.conv_states])
+        self.tstride = t([s.stride(2) for s in self.conv_states])
+        z = lambda: torch.empty(max_num_reqs, dtype=torch.int32, device=dev)
+        self.commit_lens, self.final_idx, self.bnd_idx, self.bnd_len = (
+            z(),
+            z(),
+            z(),
+            z(),
+        )
+
+    def commit(
+        self,
+        num_accepted_tokens,
+        state_indices,
+        query_start_loc,
+        request_indices=None,
+        block_table=None,
+        num_computed_tokens=None,
+        mamba_block_size=None,
+    ):
+        batch = state_indices.shape[0]
+        if batch == 0:
+            return
+        _require(state_indices.ndim == 1, "state indices must be one-dimensional")
+        _require(
+            batch <= self.commit_lens.shape[0],
+            "PLE commit batch exceeds its plan capacity",
+        )
+        _require(
+            query_start_loc.ndim == 1 and query_start_loc.shape[0] == batch + 1,
+            "PLE commit metadata is incompatible",
+        )
+        _require(
+            request_indices is None or request_indices.shape[0] >= batch,
+            "PLE request mapping is too short",
+        )
+        _require(
+            num_accepted_tokens.ndim == 1
+            and (request_indices is not None or num_accepted_tokens.shape[0] >= batch),
+            "PLE accepted-token counts are too short",
+        )
+        _align = (block_table, num_computed_tokens, mamba_block_size)
+        _require(
+            all(x is None for x in _align) or all(x is not None for x in _align),
+            "PLE align metadata is incomplete",
+        )
+        _require(
+            mamba_block_size is None or mamba_block_size >= self.spec_query_len,
+            "PLE align block size must cover one speculative window",
+        )
+        _require(
+            block_table is None or block_table.ndim == 2,
+            "PLE block table must be two-dimensional",
+        )
+        _dev = self.conv_states[0].device
+        _require(
+            all(
+                t_ is None or t_.device == _dev
+                for t_ in (
+                    num_accepted_tokens,
+                    state_indices,
+                    query_start_loc,
+                    request_indices,
+                    block_table,
+                    num_computed_tokens,
+                )
+            ),
+            "PLE commit inputs must be on the same device",
+        )
+        bt = (0, 0) if block_table is None else block_table.stride()
+        prepare_commit_plan_kernel[(batch,)](
+            num_accepted_tokens,
+            request_indices,
+            state_indices,
+            query_start_loc,
+            block_table,
+            num_computed_tokens,
+            self.commit_lens,
+            self.final_idx,
+            self.bnd_idx,
+            self.bnd_len,
+            NULL_BLOCK_ID,
+            mamba_block_size or 1,
+            block_table.shape[1] if block_table is not None else 1,
+            num_accepted_tokens.stride(0),
+            request_indices.stride(0) if request_indices is not None else 0,
+            state_indices.stride(0),
+            query_start_loc.stride(0),
+            bt[0],
+            bt[1],
+            0 if num_computed_tokens is None else num_computed_tokens.stride(0),
+            SPEC_QUERY_LEN=self.spec_query_len,
+            num_warps=1,
+        )
+        compact_conv_state_kernel[
+            (triton.cdiv(self.conv_dim, 256), batch, len(self.conv_states))
+        ](
+            self.conv_states[0],
+            self.base,
+            self.bstride,
+            self.dstride,
+            self.tstride,
+            state_indices,
+            self.commit_lens,
+            self.final_idx,
+            self.bnd_idx,
+            self.bnd_len,
+            NULL_BLOCK_ID,
+            self.conv_dim,
+            self.hist,
+            state_indices.stride(0),
+            BLOCK_D=256,
+            BLOCK_HISTORY=triton.next_power_of_2(self.hist),
+            ALIGN_MODE=block_table is not None,
+            num_warps=4,
+        )
+
+
+@dataclass
+class PleRecoverSSMCommitMetadata:
+    state_indices: torch.Tensor  # [n]
+    query_start_loc: torch.Tensor  # [n + 1]
+    request_indices: torch.Tensor | None
+    block_table: torch.Tensor | None
+    num_computed_tokens: torch.Tensor | None
+    block_size: int | None
+
+
+@dataclass
+class PleRecoverSSMMetadata(PleShortConvAttentionMetadata, RecoverSSMMetadata):
+    recoverssm_commit: PleRecoverSSMCommitMetadata | None = None
+    recoverssm_context: Any = field(default=None, repr=False, compare=False)
+
+    def commit_recoverssm_state(self, num_accepted_tokens):
+        c = self.recoverssm_commit
+        if c is None:
+            return None
+        n = self.num_spec_decodes
+        self.recoverssm_context.commit(
+            num_accepted_tokens,
+            c.state_indices[:n],
+            c.query_start_loc[: n + 1],
+            request_indices=c.request_indices,
+            block_table=c.block_table,
+            num_computed_tokens=c.num_computed_tokens,
+            mamba_block_size=c.block_size,
+        )
+        if c.block_table is None:
+            return None
+        assert c.block_size is not None and c.num_computed_tokens is not None
+        return RecoverSSMPostprocessMetadata(
+            num_spec_decodes=n,
+            request_indices=c.request_indices,
+            block_table=c.block_table,
+            num_computed_tokens=c.num_computed_tokens,
+            block_size=c.block_size,
+        )
+
+
+class PleRecoverSSMMetadataBuilder(PleShortConvAttentionMetadataBuilder):
+    # Same as the GDN RecoverSSM builder: PIECEWISE only for now.
+    _cudagraph_support = AttentionCGSupport.NEVER
+
+    def __init__(self, kv_cache_spec, layer_names, vllm_config, device) -> None:
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+        self.use_full_cuda_graph = False
+        self._ones = torch.ones(
+            vllm_config.scheduler_config.max_num_seqs, dtype=torch.int32, device=device
+        )
+        self._ctx: _PleConvCommit | None = None
+        self._logged = False
+
+    def _get_ctx(self):
+        if self._ctx is None:
+            fc = self.vllm_config.compilation_config.static_forward_context
+            convs = [fc[name].kv_cache[0] for name in self.layer_names]
+            ctx = _PleConvCommit(
+                convs, 1 + self.num_spec, self.vllm_config.scheduler_config.max_num_seqs
+            )
+            self._ctx = ctx
+            logger.info(
+                "PLE RecoverSSM commit context: %d layers, history %d",
+                len(convs),
+                ctx.hist,
+            )
+        return self._ctx
+
+    def build(
+        self,
+        common_prefix_len,
+        common_attn_metadata,
+        fast_build=False,
+        *,
+        num_accepted_tokens=None,
+        num_decode_draft_tokens_cpu=None,
+        **kwargs,
+    ):  # type: ignore[override]
+        m = common_attn_metadata
+        rows = None
+        if self.use_spec_decode and num_decode_draft_tokens_cpu is not None:
+            if m.is_prefilling is None or m.is_prefilling.device.type != "cpu":
+                raise ValueError("PLE RecoverSSM needs the CPU is_prefilling mask")
+            # the same classification as the GDN builder, so both commits map the same
+            # requests
+            num_decode_draft_tokens_cpu, rows = recoverssm_spec_rows(
+                m.is_prefilling,
+                m.query_start_loc_cpu,
+                num_decode_draft_tokens_cpu,
+                m.num_reqs,
+            )
+        meta = super().build(
+            common_prefix_len,
+            m,
+            fast_build,
+            num_accepted_tokens=num_accepted_tokens,
+            num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu,
+            **kwargs,
+        )
+        if (
+            not isinstance(meta, PleShortConvAttentionMetadata)
+            or meta.num_spec_decodes == 0
+        ):
+            return meta
+        base = {f.name: getattr(meta, f.name) for f in fields(meta)}
+        n = meta.num_spec_decodes
+        base["num_accepted_tokens"] = self._ones[:n]
+        if rows is None:
+            raise ValueError("PLE RecoverSSM: spec rows without draft counts")
+        req_idx = recoverssm_request_indices(rows, n, m.query_start_loc.device)
+        align = self.vllm_config.cache_config.mamba_cache_mode == "align"
+        commit = PleRecoverSSMCommitMetadata(
+            state_indices=meta.spec_state_indices_tensor,
+            query_start_loc=meta.spec_query_start_loc,
+            request_indices=req_idx,
+            block_table=m.block_table_tensor if align else None,
+            num_computed_tokens=m.compute_num_computed_tokens() if align else None,
+            block_size=self.kv_cache_spec.block_size if align else None,
+        )
+        if not self._logged:
+            self._logged = True
+            logger.info("PLE RecoverSSM path taken: %d spec rows, align=%s", n, align)
+        return PleRecoverSSMMetadata(
+            **base, recoverssm_commit=commit, recoverssm_context=self._get_ctx()
+        )
+
+
+class PleRecoverSSMAttentionBackend(PleShortConvAttentionBackend):
+    @staticmethod
+    def get_name() -> str:
+        return "PLE_SHORT_CONV_RECOVERSSM"
+
+    @staticmethod
+    def get_builder_cls() -> type[PleRecoverSSMMetadataBuilder]:
+        return PleRecoverSSMMetadataBuilder

vllm/v1/attention/backends/recoverssm_metadata.py

@@ -24,3 +24,41 @@ class RecoverSSMMetadata(abc.ABC):
         self, num_accepted_tokens: torch.Tensor
     ) -> RecoverSSMPostprocessMetadata | None:
         raise NotImplementedError
+
+
+def recoverssm_spec_rows(
+    is_prefilling_cpu: torch.Tensor,
+    query_start_loc_cpu: torch.Tensor,
+    num_decode_draft_tokens_cpu: torch.Tensor,
+    num_reqs: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Which batch rows take the RecoverSSM spec path, shared by the RecoverSSM metadata
+    builders (GDN, PLE).
+
+    Every active decode row does (a draft-less step is a window of one token), plus
+    every row that has drafts. Only the first ``num_reqs`` rows are real; padding rows
+    never take the spec path. Returns the draft-count vector to hand to the base builder
+    (1 for spec rows, -1 otherwise, padding included) and the spec rows in batch order,
+    which is also the order of the base builder's spec state indices."""
+    query_lens = query_start_loc_cpu[: num_reqs + 1].diff()
+    active_decode = (~is_prefilling_cpu[:num_reqs]) & (query_lens > 0)
+    spec_mask = (num_decode_draft_tokens_cpu[:num_reqs] >= 0) | active_decode
+    drafts = torch.full_like(num_decode_draft_tokens_cpu, -1)
+    drafts[:num_reqs] = torch.where(spec_mask, 1, -1).to(drafts.dtype)
+    return drafts, spec_mask.nonzero().flatten()
+
+
+def recoverssm_request_indices(
+    rows: torch.Tensor, num_spec_decodes: int, device: torch.device
+) -> torch.Tensor | None:
+    """None when the spec rows are exactly 0..n-1, else their batch rows (the commit
+    plan's request mapping). Raises if the base builder counted a different number of
+    spec rows than the classification above."""
+    if rows.numel() != num_spec_decodes:
+        raise ValueError(
+            f"RecoverSSM: {rows.numel()} spec rows classified, but the builder "
+            f"produced {num_spec_decodes}"
+        )
+    if num_spec_decodes == 0 or int(rows[-1]) == num_spec_decodes - 1:
+        return None
+    return rows.to(torch.int32).to(device, non_blocking=True)

vllm/v1/core/single_type_kv_cache_manager.py

@@ -1509,6 +1509,15 @@ class MambaManager(SingleTypeKVCacheManager):
             return computed_blocks, hit_length
 
         max_num_blocks = max_length // block_size
+        if drop_eagle_block and max_num_blocks > 0:
+            # EAGLE/MTP: drop the final matched page. Its recurrent-state snapshot
+            # may have been taken over draft tokens that verification later rejects,
+            # so the state must be restored from the previous page boundary and
+            # recomputed forward. FullAttentionManager (see find_longest_cache_hit
+            # above) pops the last block; Mamba keeps only the rightmost real block
+            # with the prefix null-padded, so a literal pop would delete the state
+            # block itself -- lower the search ceiling by one page instead.
+            max_num_blocks -= 1
         # Search from right to left and early stop when a match is found.
         for i in range(max_num_blocks - 1, -1, -1):
             if cached_block := block_pool.get_cached_block(

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -350,7 +350,7 @@ class CudaGraphManager:
 
             # recoverSSM cannot capture a dummy query wider than its workspace.
             if mixed_mode and (
-                not self.vllm_config.cache_config.use_kda_recoverssm
+                not self.vllm_config.cache_config.use_recoverssm
                 or num_tokens <= max_decode_tokens
             ):
                 # for PIECEWISE graphs there is no limit on requests when replaying

vllm/v1/worker/gpu/model_states/mamba_hybrid.py

@@ -92,7 +92,7 @@ class MambaHybridModelState(DefaultModelState):
         # columns and the running state_idx are kept GPU-resident.
         self._align_mode = self.cache_config.mamba_cache_mode == "align"
         self.recoverssm = (
-            RecoverSSMState() if self.cache_config.use_kda_recoverssm else None
+            RecoverSSMState() if self.cache_config.use_recoverssm else None
         )
         if self._align_mode:
             self._mamba_state_idx_gpu = torch.zeros(

vllm/v1/worker/gpu/model_states/recoverssm.py

@@ -91,6 +91,8 @@ def _postprocess_recoverssm_align_kernel(
         return
     num_sampled = tl.load(num_sampled_ptr + batch_idx)
     num_computed = tl.load(num_computed_ptr + batch_idx)
+    # Match the commit plan: the running state lives in the block holding the
+    # last computed token, so the next step migrates it across a boundary.
     tl.store(
         state_idx_ptr + req_state_idx,
         tl.minimum(