Remove obsolete sgl-kernel legacy paths (#21528)

python/sglang/jit_kernel/benchmark/bench_cast.py

@@ -1,7 +1,6 @@
 import torch
 import triton
 import triton.testing
-from sgl_kernel import downcast_fp8 as downcast_fp8_aot
 
 from sglang.jit_kernel.benchmark.utils import (
     DEFAULT_DEVICE,
@@ -31,9 +30,9 @@ HEAD_DIM_LIST = get_benchmark_range(
 
 CONFIGS = [(sl, h, d, sl * 2) for sl in SL_LIST for h, d in HEAD_DIM_LIST]
 
-LINE_VALS = ["aot", "jit"]
+LINE_VALS = ["jit"]
-LINE_NAMES = ["AOT (sgl-kernel)", "JIT (cast.cuh, 256 threads, 2D grid)"]
+LINE_NAMES = ["JIT (cast.cuh, 256 threads, 2D grid)"]
-STYLES = [("blue", "--"), ("orange", "-")]
+STYLES = [("orange", "-")]
 
 
 # ── Perf report ────────────────────────────────────────────────────────────────
@@ -48,7 +47,7 @@ STYLES = [("blue", "--"), ("orange", "-")]
         line_names=LINE_NAMES,
         styles=STYLES,
         ylabel="us",
-        plot_name="downcast-fp8-aot-vs-jit",
+        plot_name="downcast-fp8-jit",
         args={},
     )
 )
@@ -61,10 +60,7 @@ def benchmark(input_sl, head, dim, out_sl, provider):
     v_scale = torch.tensor([1.0], dtype=torch.float32, device=DEVICE)
     loc = torch.arange(input_sl, dtype=torch.int64, device=DEVICE)
 
-    if provider == "aot":
+    fn = lambda: downcast_fp8_jit(k, v, k_out, v_out, k_scale, v_scale, loc)
-        fn = lambda: downcast_fp8_aot(k, v, k_out, v_out, k_scale, v_scale, loc)
-    else:
-        fn = lambda: downcast_fp8_jit(k, v, k_out, v_out, k_scale, v_scale, loc)
 
     return run_benchmark(fn)
 
@@ -84,26 +80,19 @@ def _report_bandwidth(input_sl, head, dim, dtype):
     v_scale = torch.tensor([1.0], dtype=torch.float32, device=DEVICE)
     loc = torch.arange(input_sl, dtype=torch.int64, device=DEVICE)
 
-    aot_fn = lambda: downcast_fp8_aot(k, v, k_out, v_out, k_scale, v_scale, loc)
     jit_fn = lambda: downcast_fp8_jit(k, v, k_out, v_out, k_scale, v_scale, loc)
 
-    aot_ms, _, _ = triton.testing.do_bench(aot_fn, quantiles=[0.5, 0.2, 0.8])
     jit_ms, _, _ = triton.testing.do_bench(jit_fn, quantiles=[0.5, 0.2, 0.8])
 
     def fmt(ms):
         return f"{ms*1000:6.2f}us {total_bytes/(ms*1e-3)/1e9:6.0f}GB/s"
 
-    print(
+    print(f"  sl={input_sl:5d}  h={head:2d}  d={dim:4d}" f"  |  jit {fmt(jit_ms)}")
-        f"  sl={input_sl:5d}  h={head:2d}  d={dim:4d}"
-        f"  |  aot {fmt(aot_ms)}"
-        f"  |  jit {fmt(jit_ms)}"
-        f"  |  speedup {aot_ms/jit_ms:.2f}x"
-    )
 
 
 def report_bandwidth():
     print(f"\n{'='*95}")
-    print("  AOT (sgl-kernel) vs JIT (cast.cuh, 256 threads, 2D grid)")
+    print("  JIT (cast.cuh, 256 threads, 2D grid)")
     print(f"  dtype={DTYPE}, device={DEVICE}")
     print(f"{'='*95}")
     for sl in [64, 256, 1024, 2048]:

python/sglang/jit_kernel/benchmark/bench_renorm.py

@@ -82,31 +82,6 @@ def torch_top_p_renorm_probs(probs, top_p, eps=1e-5):
         return renorm_probs
 
 
-def torch_top_k_mask_logits(logits, top_k):
-    """Vectorized PyTorch implementation of top-k logits masking."""
-    batch_size, vocab_size = logits.shape
-
-    # Handle scalar or tensor k
-    if isinstance(top_k, int):
-        k_val = min(max(top_k, 1), vocab_size)
-        # Get top-k indices for all batches at once
-        _, topk_indices = torch.topk(logits, k_val, dim=1, largest=True)
-
-        # Create masked logits: start with -inf everywhere
-        masked_logits = torch.full_like(logits, float("-inf"))
-        # Scatter the top-k values back
-        masked_logits.scatter_(1, topk_indices, logits.gather(1, topk_indices))
-    else:
-        # Variable k per batch - need to handle separately
-        masked_logits = torch.full_like(logits, float("-inf"))
-        for i in range(batch_size):
-            k_val = min(max(top_k[i].item(), 1), vocab_size)
-            _, topk_indices = torch.topk(logits[i], k_val, largest=True)
-            masked_logits[i, topk_indices] = logits[i, topk_indices]
-
-    return masked_logits
-
-
 def calculate_diff_top_k_renorm(batch_size, vocab_size, k):
     """Compare Torch reference and SGLang kernel for top-k renorm correctness."""
     torch.manual_seed(42)
@@ -139,20 +114,6 @@ def calculate_diff_top_p_renorm(batch_size, vocab_size, p):
     torch.testing.assert_close(torch_output, sglang_output, rtol=1e-3, atol=1e-3)
 
 
-def calculate_diff_top_k_mask(batch_size, vocab_size, k):
-    """Compare Torch reference and SGLang kernel for top-k mask correctness."""
-    torch.manual_seed(42)
-    device = torch.device("cuda")
-
-    logits = torch.randn(batch_size, vocab_size, device=device) * 5
-    top_k_tensor = torch.full((batch_size,), k, device=device, dtype=torch.int32)
-
-    torch_output = torch_top_k_mask_logits(logits, top_k_tensor)
-    sglang_output = sgl_kernel.top_k_mask_logits(logits, top_k_tensor)
-
-    torch.testing.assert_close(torch_output, sglang_output, rtol=1e-3, atol=1e-3)
-
-
 # Parameter space - simplified for CI
 if is_in_ci():
     batch_size_range = [16]
@@ -231,38 +192,6 @@ def benchmark_top_p_renorm(batch_size, vocab_size, p, provider):
     return run_benchmark_no_cudagraph(fn)
 
 
-@triton.testing.perf_report(
-    triton.testing.Benchmark(
-        x_names=["batch_size", "vocab_size", "k"],
-        x_vals=configs_k,
-        line_arg="provider",
-        line_vals=["torch", "sglang"],
-        line_names=["Torch Reference", "SGL Kernel"],
-        styles=[("red", "-"), ("orange", "-")],
-        ylabel="us",
-        plot_name="top-k-mask-logits-performance",
-        args={},
-    )
-)
-def benchmark_top_k_mask(batch_size, vocab_size, k, provider):
-    # Skip invalid configurations
-    if k >= vocab_size:
-        return float("nan"), float("nan"), float("nan")
-
-    torch.manual_seed(42)
-    device = torch.device("cuda")
-
-    logits = torch.randn(batch_size, vocab_size, device=device) * 5
-    top_k_tensor = torch.full((batch_size,), k, device=device, dtype=torch.int32)
-
-    if provider == "torch":
-        fn = lambda: torch_top_k_mask_logits(logits.clone(), top_k_tensor)
-    elif provider == "sglang":
-        fn = lambda: sgl_kernel.top_k_mask_logits(logits.clone(), top_k_tensor)
-
-    return run_benchmark_no_cudagraph(fn)
-
-
 if __name__ == "__main__":
     print("=" * 60)
     print("Running correctness checks...")
@@ -291,15 +220,6 @@ if __name__ == "__main__":
         batch_size, vocab_size, p = cfg
         print(f"  ✓ Passed: batch_size={batch_size}, vocab_size={vocab_size}, p={p}")
 
-    print("\n3. Testing top_k_mask_logits...")
-    for cfg in test_configs_k:
-        batch_size, vocab_size, k = cfg
-        if k < vocab_size:  # Skip invalid configs
-            calculate_diff_top_k_mask(batch_size, vocab_size, k)
-            print(
-                f"  ✓ Passed: batch_size={batch_size}, vocab_size={vocab_size}, k={k}"
-            )
-
     print("\n" + "=" * 60)
     print("All correctness checks passed!")
     print("=" * 60)
@@ -314,9 +234,6 @@ if __name__ == "__main__":
     print("\n2. Benchmarking top_p_renorm_probs...")
     benchmark_top_p_renorm.run(print_data=True)
 
-    print("\n3. Benchmarking top_k_mask_logits...")
-    benchmark_top_k_mask.run(print_data=True)
-
     print("\n" + "=" * 60)
     print("Benchmarking complete!")
     print("=" * 60)

python/sglang/jit_kernel/tests/test_renorm.py

@@ -82,44 +82,5 @@ def test_top_p_renorm_probs(batch_size, vocab_size, p):
     )
 
 
-@pytest.mark.parametrize("batch_size", [1, 99, 989])
-@pytest.mark.parametrize("vocab_size", [111, 32000, 128256])
-@pytest.mark.parametrize("k", [10, 100, 500])
-@pytest.mark.parametrize("neginf_input", [False, True])
-def test_top_k_mask_logits(batch_size, vocab_size, k, neginf_input):
-    """Test top_k_mask_logits kernel for correctness.
-
-    This test validates that the kernel correctly:
-    1. Identifies the top-k logits
-    2. Masks non-top-k values to -inf
-    3. Preserves the top-k values
-    4. Handles negative infinity inputs gracefully
-
-    The test verifies correctness by comparing softmax(top_k_mask_logits(logits))
-    with top_k_renorm_prob(probs), which should be equivalent.
-    """
-    if k > vocab_size:
-        pytest.skip("k should be less than vocab_size")
-    torch.manual_seed(42)
-    logits = torch.randn(batch_size, vocab_size, device="cuda:0") * 5
-    if neginf_input:
-        # Randomly assign some logits to -inf to test edge cases
-        num_neginf = torch.randint(1, vocab_size * batch_size, (1,)).item()
-        idxs = torch.randperm(batch_size * vocab_size, device="cuda:0")[:num_neginf]
-        logits[idxs // vocab_size, idxs % vocab_size] = -float("inf")
-
-    probs = torch.softmax(logits, dim=-1)
-    masked_logits = sgl_kernel.top_k_mask_logits(logits, k)
-    renormed_probs = torch.softmax(masked_logits, dim=-1)
-    renormed_probs_ref = sgl_kernel.top_k_renorm_prob(probs, k)
-
-    torch.testing.assert_close(
-        renormed_probs,
-        renormed_probs_ref,
-        rtol=1e-3,
-        atol=1e-3,
-    )
-
-
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))

sgl-kernel/CMakeLists.txt

@@ -260,13 +260,11 @@ endif()
 set(SOURCES
     "csrc/allreduce/custom_all_reduce.cu"
     "csrc/allreduce/mscclpp_allreduce.cu"
-    "csrc/attention/cascade.cu"
     "csrc/attention/cutlass_mla_kernel.cu"
     "csrc/attention/merge_attn_states.cu"
     "csrc/attention/vertical_slash_index.cu"
     "csrc/common_extension.cc"
     "csrc/elementwise/activation.cu"
-    "csrc/elementwise/cast.cu"
     "csrc/elementwise/concat_mla.cu"
     "csrc/elementwise/copy.cu"
     "csrc/elementwise/fused_add_rms_norm_kernel.cu"

sgl-kernel/csrc/attention/cascade.cu

@@ -1,55 +0,0 @@
-// Adapted from
-// https://github.com/flashinfer-ai/flashinfer/blob/55576c626421b5ee7e7ebe74afd26465c8ae863f/csrc/cascade.cu
-
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-
-#include <flashinfer/attention/cascade.cuh>
-
-#include "pytorch_extension_utils.h"
-
-using namespace flashinfer;
-
-void merge_state(
-    at::Tensor v_a, at::Tensor s_a, at::Tensor v_b, at::Tensor s_b, at::Tensor v_merged, at::Tensor s_merged) {
-  CHECK_INPUT(v_a);
-  CHECK_INPUT(s_a);
-  CHECK_INPUT(v_b);
-  CHECK_INPUT(s_b);
-  auto device = v_a.device();
-  CHECK_EQ(s_a.device(), device);
-  CHECK_EQ(v_b.device(), device);
-  CHECK_EQ(s_b.device(), device);
-  CHECK_DIM(3, v_a);
-  CHECK_DIM(2, s_a);
-  CHECK_DIM(3, v_b);
-  CHECK_DIM(2, s_b);
-  CHECK_SHAPE(v_a, v_b);
-  CHECK_SHAPE(s_a, s_b);
-  CHECK_EQ(v_a.size(0), s_a.size(0));
-  CHECK_EQ(v_a.size(1), s_b.size(1));
-  unsigned int seq_len = v_a.size(0);
-  unsigned int num_heads = v_a.size(1);
-  unsigned int head_dim = v_a.size(2);
-
-  const c10::cuda::OptionalCUDAGuard device_guard(v_a.device());
-  auto stream = at::cuda::getCurrentCUDAStream();
-
-  bool success = DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(v_a.scalar_type(), c_type, [&] {
-    cudaError_t status = MergeState(
-        static_cast<c_type*>(v_a.data_ptr()),
-        static_cast<float*>(s_a.data_ptr()),
-        static_cast<c_type*>(v_b.data_ptr()),
-        static_cast<float*>(s_b.data_ptr()),
-        static_cast<c_type*>(v_merged.data_ptr()),
-        static_cast<float*>(s_merged.data_ptr()),
-        seq_len,
-        num_heads,
-        head_dim,
-        stream);
-    TORCH_CHECK(status == cudaSuccess, "MergeState kernel launch failed: ", cudaGetErrorString(status));
-    return true;
-  });
-
-  TORCH_CHECK(success, "MergeState kernel launch failed: unsupported data type");
-}

sgl-kernel/csrc/common_extension.cc

@@ -50,8 +50,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
   /*
    * From csrc/attention
    */
-  m.def("merge_state(Tensor v_a, Tensor s_a, Tensor v_b, Tensor s_b, Tensor! v_merged, Tensor! s_merged) -> ()");
-  m.impl("merge_state", torch::kCUDA, &merge_state);
   m.def("merge_state_v2(Tensor v_a, Tensor s_a, Tensor v_b, Tensor s_b, Tensor! v_merged, Tensor! s_merged) -> ()");
   m.impl("merge_state_v2", torch::kCUDA, &merge_state_v2);
   m.def(
@@ -90,11 +88,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
       "                 Tensor cos_sin_cache, bool is_neox) -> ()");
   m.impl("rotary_embedding", torch::kCUDA, &rotary_embedding);
 
-  m.def(
-      "downcast_fp8(Tensor k, Tensor v, Tensor k_out, Tensor v_out, Tensor k_scale, Tensor v_scale, Tensor loc, "
-      "int mult, int offset) -> ()");
-  m.impl("downcast_fp8", torch::kCUDA, &downcast_fp8);
-
   m.def("copy_to_gpu_no_ce(Tensor input, Tensor! output) -> ()");
   m.impl("copy_to_gpu_no_ce", torch::kCUDA, &copy_to_gpu_no_ce);
   m.def("concat_mla_k(Tensor! k, Tensor k_nope, Tensor k_rope) -> ()");
@@ -364,9 +357,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
   m.def("top_p_renorm_probs(Tensor probs, Tensor! renorm_probs, Tensor? maybe_top_p_arr, float top_p_val) -> ()");
   m.impl("top_p_renorm_probs", torch::kCUDA, &top_p_renorm_probs);
 
-  m.def("top_k_mask_logits(Tensor logits, Tensor mask_logits, Tensor? maybe_top_k_arr, int top_k_val) -> ()");
-  m.impl("top_k_mask_logits", torch::kCUDA, &top_k_mask_logits);
-
   /*
    * From Sparse Flash Attention
    */

sgl-kernel/csrc/common_extension_musa.cc

@@ -43,9 +43,6 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
       "top_k_top_p_sampling_from_probs(Tensor probs, Tensor output, Tensor? maybe_indices, Tensor? maybe_top_k_arr, "
       "float top_k_val, Tensor? maybe_top_p_arr, float top_p_val, bool deterministic, Generator? gen) -> ()");
   m.impl("top_k_top_p_sampling_from_probs", torch::kMUSA, &top_k_top_p_sampling_from_probs);
-
-  m.def("top_k_mask_logits(Tensor logits, Tensor mask_logits, Tensor? maybe_top_k_arr, int top_k_val) -> ()");
-  m.impl("top_k_mask_logits", torch::kMUSA, &top_k_mask_logits);
 }
 
 REGISTER_EXTENSION(common_ops)

sgl-kernel/csrc/elementwise/cast.cu

@@ -1,172 +0,0 @@
-#include <ATen/cuda/CUDAContext.h>
-
-#include "utils.h"
-
-template <typename T>
-struct ConvertToFP8 {
-  static __device__ __nv_fp8_storage_t convert_to_fp8(T value) {
-    return 0;
-  }
-};
-
-template <>
-struct ConvertToFP8<__nv_bfloat16> {
-  static __device__ __nv_fp8_storage_t convert_to_fp8(__nv_bfloat16 value) {
-    return __nv_cvt_bfloat16raw_to_fp8(value, __NV_SATFINITE, __NV_E4M3);
-  }
-};
-
-template <>
-struct ConvertToFP8<half> {
-  static __device__ __nv_fp8_storage_t convert_to_fp8(half value) {
-    return __nv_cvt_halfraw_to_fp8(value, __NV_SATFINITE, __NV_E4M3);
-  }
-};
-
-template <typename T>
-struct ConvertFromFloat {
-  static __device__ T convert_from_float(float value) {
-    return 0;
-  }
-};
-
-template <>
-struct ConvertFromFloat<__nv_bfloat16> {
-  static __device__ __nv_bfloat16 convert_from_float(float value) {
-    return __float2bfloat16(value);
-  }
-};
-
-template <>
-struct ConvertFromFloat<half> {
-  static __device__ half convert_from_float(float value) {
-    return __float2half(value);
-  }
-};
-
-template <typename T>
-__global__ void fused_downcast_kernel(
-    const T* cache_k,
-    const T* cache_v,
-    const float* k_scale,
-    const float* v_scale,
-    __nv_fp8_storage_t* output_k,
-    __nv_fp8_storage_t* output_v,
-    const int input_sl,
-    const int head,
-    const int dim,
-    const T max_fp8,
-    const T min_fp8,
-    const int64_t mult,
-    const int64_t offset,
-    const int64_t* loc) {
-  // TODO: change name
-  int token_idx = blockIdx.x;
-  int thread_idx = threadIdx.x;
-  int total_threads = blockDim.x;
-
-  T k_scale_val = ConvertFromFloat<T>::convert_from_float(k_scale[0]);
-  T v_scale_val = ConvertFromFloat<T>::convert_from_float(v_scale[0]);
-
-  T k_scale_inv = static_cast<T>(1.f) / k_scale_val;
-  T v_scale_inv = static_cast<T>(1.f) / v_scale_val;
-
-  auto clamp = [&](T val) { return val > max_fp8 ? max_fp8 : (min_fp8 > val ? min_fp8 : val); };
-
-  if (token_idx < input_sl) {
-    int out_seq_idx = loc[token_idx];
-
-#pragma unroll
-    for (int i = thread_idx; i < head * dim; i += total_threads) {
-      int in_idx = token_idx * head * dim + i;
-      int out_idx = (out_seq_idx * mult + offset) * head * dim + i;
-
-      T k_val = cache_k[in_idx] * k_scale_inv;
-      k_val = clamp(k_val);
-      output_k[out_idx] = ConvertToFP8<T>::convert_to_fp8(k_val);
-
-      T v_val = cache_v[in_idx] * v_scale_inv;
-      v_val = clamp(v_val);
-      output_v[out_idx] = ConvertToFP8<T>::convert_to_fp8(v_val);
-    }
-  }
-}
-
-template <typename T>
-void downcast_fp8_impl(
-    at::Tensor& k,
-    at::Tensor& v,
-    at::Tensor& k_out,
-    at::Tensor& v_out,
-    at::Tensor& k_scale,
-    at::Tensor& v_scale,
-    at::Tensor& loc,
-    int64_t mult,
-    int64_t offset,
-    cudaStream_t stream) {
-  CHECK_INPUT(k);
-  CHECK_INPUT(v);
-  CHECK_INPUT(k_out);
-  CHECK_INPUT(v_out);
-  CHECK_INPUT(k_scale);
-  CHECK_INPUT(v_scale);
-  CHECK_INPUT(loc);
-
-  int64_t input_sl = k.size(0);
-  int64_t head = k.size(1);
-  int64_t dim = k.size(2);
-
-  dim3 grid(input_sl * head);
-  int vec_size = 8;
-  dim3 block(std::min(int(dim) / vec_size, 1024));
-
-  const T max_fp8 = static_cast<T>(FP8_E4M3_MAX);
-  const T min_fp8 = static_cast<T>(-FP8_E4M3_MAX);
-
-  fused_downcast_kernel<T><<<grid, block, 0, stream>>>(
-      static_cast<const T*>(k.data_ptr()),
-      static_cast<const T*>(v.data_ptr()),
-      static_cast<const float*>(k_scale.data_ptr()),
-      static_cast<const float*>(v_scale.data_ptr()),
-      static_cast<__nv_fp8_storage_t*>(k_out.data_ptr()),
-      static_cast<__nv_fp8_storage_t*>(v_out.data_ptr()),
-      input_sl,
-      head,
-      dim,
-      max_fp8,
-      min_fp8,
-      mult,
-      offset,
-      static_cast<const int64_t*>(loc.data_ptr()));
-
-  cudaError_t status = cudaGetLastError();
-  TORCH_CHECK(status == cudaSuccess, "Kernel launch failed: " + std::string(cudaGetErrorString(status)));
-}
-
-void downcast_fp8(
-    at::Tensor& k,
-    at::Tensor& v,
-    at::Tensor& k_out,
-    at::Tensor& v_out,
-    at::Tensor& k_scale,
-    at::Tensor& v_scale,
-    at::Tensor& loc,
-    int64_t mult,
-    int64_t offset) {
-  CHECK_INPUT(k);
-  CHECK_INPUT(v);
-  CHECK_INPUT(k_out);
-  CHECK_INPUT(v_out);
-
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  switch (k.scalar_type()) {
-    case at::ScalarType::BFloat16:
-      downcast_fp8_impl<__nv_bfloat16>(k, v, k_out, v_out, k_scale, v_scale, loc, mult, offset, stream);
-      break;
-    case at::ScalarType::Half:
-      downcast_fp8_impl<__half>(k, v, k_out, v_out, k_scale, v_scale, loc, mult, offset, stream);
-      break;
-    default:
-      TORCH_CHECK(false, "Unsupported input type for downcast_fp8. Expected bfloat16 or float16.");
-  }
-}

sgl-kernel/include/sgl_kernel_ops.h

@@ -103,8 +103,6 @@ void mscclpp_allreduce(fptr_t _context, torch::Tensor& inp, torch::Tensor& out,
 /*
  * From csrc/attention
  */
-void merge_state(
-    at::Tensor v_a, at::Tensor s_a, at::Tensor v_b, at::Tensor s_b, at::Tensor v_merged, at::Tensor s_merged);
 void merge_state_v2(
     at::Tensor v_a, at::Tensor s_a, at::Tensor v_b, at::Tensor s_b, at::Tensor v_merged, at::Tensor s_merged);
 void cutlass_mla_decode(
@@ -143,17 +141,6 @@ void rotary_embedding(
     torch::Tensor& cos_sin_cache,
     bool is_neox);
 
-void downcast_fp8(
-    at::Tensor& k,
-    at::Tensor& v,
-    at::Tensor& k_out,
-    at::Tensor& v_out,
-    at::Tensor& k_scale,
-    at::Tensor& v_scale,
-    at::Tensor& loc,
-    int64_t mult,
-    int64_t offset);
-
 void copy_to_gpu_no_ce(const at::Tensor& input, at::Tensor& output);
 void concat_mla_k(torch::Tensor k, torch::Tensor k_nope, torch::Tensor k_rope);
 void concat_mla_absorb_q(at::Tensor a, at::Tensor b, at::Tensor out);
@@ -604,9 +591,6 @@ void top_k_renorm_probs(
 void top_p_renorm_probs(
     at::Tensor probs, at::Tensor renorm_probs, std::optional<at::Tensor> maybe_top_p_arr, double top_p_val);
 
-void top_k_mask_logits(
-    at::Tensor logits, at::Tensor mask_logits, std::optional<at::Tensor> maybe_top_k_arr, int64_t top_k_val);
-
 namespace flash {
 /*
  * From fa2 sparse

sgl-kernel/python/sgl_kernel/__init__.py

@@ -14,7 +14,6 @@ from sgl_kernel.allreduce import *
 from sgl_kernel.attention import (
     cutlass_mla_decode,
     cutlass_mla_get_workspace_size,
-    merge_state,
     merge_state_v2,
 )
 from sgl_kernel.cutlass_moe import cutlass_w4a8_moe_mm, get_cutlass_w4a8_moe_mm_data
@@ -22,7 +21,6 @@ from sgl_kernel.elementwise import (
     concat_mla_absorb_q,
     concat_mla_k,
     copy_to_gpu_no_ce,
-    downcast_fp8,
     fused_add_rmsnorm,
     gelu_and_mul,
     gelu_tanh_and_mul,
@@ -92,7 +90,6 @@ from sgl_kernel.quantization import (
     ggml_mul_mat_vec_a8,
 )
 from sgl_kernel.sampling import (
-    top_k_mask_logits,
     top_k_renorm_prob,
     top_p_renorm_prob,
 )
@@ -128,7 +125,6 @@ _DEBUG_EXPORT_NAMES = [
     "copy_to_gpu_no_ce",
     "cutlass_mla_decode",
     "cutlass_mla_get_workspace_size",
-    "downcast_fp8",
     "dsv3_fused_a_gemm",
     "dsv3_router_gemm",
     "es_fp8_blockwise_scaled_grouped_mm",
@@ -151,7 +147,6 @@ _DEBUG_EXPORT_NAMES = [
     "gptq_shuffle",
     "int8_scaled_mm",
     "kimi_k2_moe_fused_gate",
-    "merge_state",
     "merge_state_v2",
     "moe_align_block_size",
     "moe_fused_gate",
@@ -170,7 +165,6 @@ _DEBUG_EXPORT_NAMES = [
     "sgl_per_token_quant_fp8",
     "shuffle_rows",
     "silu_and_mul",
-    "top_k_mask_logits",
     "top_k_renorm_prob",
     "top_p_renorm_prob",
     "topk_sigmoid",

sgl-kernel/python/sgl_kernel/_fa4_interface.py

@@ -1,940 +0,0 @@
-# Adapted from https://github.com/Dao-AILab/flash-attention/blob/5d4c9537a1e0f1adcc3e4c3e11ae46fe94a18b11/flash_attn/cute/interface.py
-
-# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
-# [2025-10-14] Version in Cute-DSL, for Hopper and Blackwell. You'd need to install nvidia-cutlass-dsl==4.2.1.
-
-
-import copy
-import gc
-import logging
-import math
-import os
-from functools import lru_cache
-from typing import Callable, Optional, Tuple
-
-logger = logging.getLogger(__name__)
-
-
-import cuda.bindings.driver as cuda
-import cutlass
-import cutlass.cute as cute
-import torch
-from cutlass.cute.runtime import from_dlpack
-from flash_attn_origin.cute import utils
-from flash_attn_origin.cute.block_sparsity import (
-    BlockSparseTensorsTorch,
-    get_block_sparse_expected_shapes,
-    normalize_block_sparse_tensors,
-    to_cute_block_sparse_tensors,
-)
-from flash_attn_origin.cute.flash_fwd import FlashAttentionForwardSm90
-from flash_attn_origin.cute.flash_fwd_combine import FlashAttentionForwardCombine
-from flash_attn_origin.cute.flash_fwd_sm100 import FlashAttentionForwardSm100
-
-
-@lru_cache(maxsize=None)
-def _get_device_capability():
-    """Cached device capability check."""
-    return torch.cuda.get_device_capability()[0]
-
-
-def maybe_contiguous(x):
-    return x.contiguous() if x is not None and x.stride(-1) != 1 else x
-
-
-def _validate_tensor(t, name, expected_shape, expected_dtype, expected_device):
-    assert (
-        t.shape == expected_shape
-    ), f"{name} shape {t.shape} != expected {expected_shape}"
-    assert (
-        t.dtype == expected_dtype
-    ), f"{name} dtype {t.dtype} != expected {expected_dtype}"
-    assert (
-        t.device == expected_device
-    ), f"{name} device {t.device} != expected {expected_device}"
-    assert t.is_cuda, f"{name} must be on CUDA"
-
-
-def to_cute_tensor(t, assumed_align=16, leading_dim=-1, fully_dynamic=False):
-    """Convert torch tensor to cute tensor for TVM FFI. leading_dim=-1 defaults to t.ndim-1."""
-    tensor = from_dlpack(t.detach(), assumed_align=assumed_align, enable_tvm_ffi=True)
-    if fully_dynamic:
-        return tensor.mark_layout_dynamic()
-    if leading_dim == -1:
-        leading_dim = t.ndim - 1
-    return tensor.mark_layout_dynamic(leading_dim=leading_dim)
-
-
-torch2cute_dtype_map = {
-    torch.float16: cutlass.Float16,
-    torch.bfloat16: cutlass.BFloat16,
-    torch.float32: cutlass.Float32,
-}
-
-
-def num_splits_heuristic(total_mblocks, num_SMs, num_n_blocks, max_splits):
-    # If num_n_blocks is too small, use 1 split. For example, we never split for hdim = 128 and seqlen_k = 512.
-    if num_n_blocks <= 4:
-        return 1
-
-    # NOTE: We should revisit this heuristic after persistence is supported for split KV.
-    # Sometimes, it's ideal to over-schedule splits for better efficiency.
-    return min(num_SMs // total_mblocks, max_splits, num_n_blocks)
-
-
-def _flash_attn_fwd(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    cu_seqlens_q: Optional[torch.Tensor] = None,
-    cu_seqlens_k: Optional[torch.Tensor] = None,
-    seqused_q: Optional[torch.Tensor] = None,
-    seqused_k: Optional[torch.Tensor] = None,
-    max_seqlen_q: Optional[int] = None,
-    max_seqlen_k: Optional[int] = None,
-    page_table: Optional[torch.Tensor] = None,
-    softmax_scale: Optional[float] = None,
-    causal: bool = False,
-    softcap: Optional[float] = None,
-    window_size_left: Optional[int] = None,
-    window_size_right: Optional[int] = None,
-    learnable_sink: Optional[torch.Tensor] = None,
-    # m_block_size: int = 128,
-    # n_block_size: int = 64,
-    # num_threads: int = 128,
-    m_block_size: int = 128,
-    n_block_size: int = 128,
-    num_threads: int = 384,
-    num_splits: int = 1,
-    pack_gqa: Optional[bool] = None,
-    _compute_capability: Optional[int] = None,
-    score_mod: Optional[Callable] = None,
-    mask_mod: Optional[Callable] = None,
-    block_sparse_tensors: Optional[BlockSparseTensorsTorch] = None,
-    return_lse: bool = False,
-    out: Optional[torch.Tensor] = None,
-    lse: Optional[torch.Tensor] = None,
-    aux_tensors: Optional[list[torch.Tensor]] = None,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Forward pass for FlashAttention.
-
-    Args:
-        ...
-        score_mod: A callable that takes the attention scores and applies a modification.
-        mask_mod: A callable that takes token position information and selectively masks
-        block_sparse_tensors: A tuple of tensors used for block sparsity.
-        return_lse: Whether to return the log softmax of the attention scores. If set to True will always calculate
-        out: Optional pre-allocated output tensor. If None, will be allocated internally.
-        lse: Optional pre-allocated log-sum-exp tensor. If None, will be allocated when needed.
-        aux_tensors: Some score_mods will want to read from global aux_tensors. This is how we thread them through to the inner kernel.
-    """
-    q, k, v = [maybe_contiguous(t) for t in (q, k, v)]
-    num_head, head_dim = q.shape[-2:]
-    if cu_seqlens_q is None:
-        batch_size, seqlen_q = q.shape[:2]
-        total_q = batch_size * seqlen_q
-    else:
-        batch_size = cu_seqlens_q.shape[0] - 1
-        seqlen_q = None
-        total_q = q.shape[0]
-    if page_table is not None:
-        assert cu_seqlens_k is None, "page_table is not supported with cu_seqlens_k"
-        assert page_table.dtype == torch.int32, "page_table must be int32"
-        assert (
-            page_table.stride(-1) == 1
-        ), "page_table must be contiguous in the last dimension"
-        max_num_pages_per_seq = page_table.shape[1]
-        assert page_table.shape == (batch_size, max_num_pages_per_seq)
-        num_pages, page_size = k.shape[:2]
-        seqlen_k = num_pages * page_size
-    else:
-        num_pages, page_size = None, None
-        seqlen_k = k.shape[-3]
-    num_head_kv = k.shape[-2]
-    head_dim_v = v.shape[-1]
-    if cu_seqlens_k is None:
-        if page_table is None:
-            assert k.shape == (batch_size, seqlen_k, num_head_kv, head_dim)
-            assert v.shape == (batch_size, seqlen_k, num_head_kv, head_dim_v)
-        else:
-            assert k.shape == (num_pages, page_size, num_head_kv, head_dim)
-            assert v.shape == (num_pages, page_size, num_head_kv, head_dim_v)
-    else:
-        assert k.shape == (seqlen_k, num_head_kv, head_dim)
-        assert v.shape == (seqlen_k, num_head_kv, head_dim_v)
-        assert cu_seqlens_k.shape == (
-            batch_size + 1,
-        ), "cu_seqlens_k must have shape (batch_size + 1,)"
-
-    if cu_seqlens_q is not None:
-        assert cu_seqlens_q.shape == (
-            batch_size + 1,
-        ), "cu_seqlens_q must have shape (batch_size + 1,)"
-    assert seqused_q is None or seqused_q.shape == (
-        batch_size,
-    ), "seqused_q must have shape (batch_size,)"
-    assert seqused_k is None or seqused_k.shape == (
-        batch_size,
-    ), "seqused_k must have shape (batch_size,)"
-    assert q.dtype in [
-        torch.float16,
-        torch.bfloat16,
-    ], "inputs must be float16 or bfloat16"
-    assert q.dtype == k.dtype == v.dtype, "inputs must have the same dtype"
-    for t in [cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k]:
-        if t is not None:
-            assert (
-                t.dtype == torch.int32
-            ), "cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k must be int32"
-            assert (
-                t.stride(0) == 1
-            ), "cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k must be contiguous"
-    if learnable_sink is not None:
-        assert learnable_sink.shape == (num_head,)
-        assert learnable_sink.dtype == torch.bfloat16, "learnable_sink must be bfloat16"
-
-    assert all(
-        t is None or t.is_cuda
-        for t in (
-            q,
-            k,
-            v,
-            cu_seqlens_q,
-            cu_seqlens_k,
-            seqused_q,
-            seqused_k,
-            page_table,
-            learnable_sink,
-        )
-    ), "inputs must be on CUDA device"
-    assert num_head % num_head_kv == 0, "num_head must be divisible by num_head_kv"
-    assert head_dim <= 256, "head_dim must be less than or equal to 256"
-    alignment = 16 // q.element_size()
-    assert head_dim % alignment == 0, f"head_dim must be divisible by {alignment}"
-    assert head_dim_v % alignment == 0, f"head_dim_v must be divisible by {alignment}"
-    if softmax_scale is None:
-        softmax_scale = 1.0 / math.sqrt(head_dim)
-    if softcap == 0.0:
-        softcap = None
-    qhead_per_kvhead = num_head // num_head_kv
-    if pack_gqa is None:
-        pack_gqa = qhead_per_kvhead > 1
-
-    out_torch_dtype = q.dtype
-    device = q.device
-    q_batch_seqlen_shape = (
-        (batch_size, seqlen_q) if cu_seqlens_q is None else (total_q,)
-    )
-    lse_shape = (
-        (batch_size, num_head, seqlen_q)
-        if cu_seqlens_q is None
-        else (num_head, total_q)
-    )
-    requires_grad = q.requires_grad or k.requires_grad or v.requires_grad
-
-    if out is None:
-        out = torch.empty(
-            *q_batch_seqlen_shape,
-            num_head,
-            head_dim_v,
-            dtype=out_torch_dtype,
-            device=device,
-        )
-    else:
-        _validate_tensor(
-            out,
-            "out",
-            (*q_batch_seqlen_shape, num_head, head_dim_v),
-            out_torch_dtype,
-            device,
-        )
-
-    if lse is None:
-        lse = (
-            torch.empty(lse_shape, dtype=torch.float32, device=device)
-            if requires_grad or return_lse
-            else None
-        )
-    elif lse is not None:
-        _validate_tensor(lse, "lse", lse_shape, torch.float32, device)
-
-    dtype = torch2cute_dtype_map[q.dtype]
-    compute_capability = (
-        _get_device_capability() if _compute_capability is None else _compute_capability
-    )
-
-    assert compute_capability in [
-        9,
-        10,
-        11,
-    ], "Unsupported compute capability. Supported: 9.x, 10.x, 11.x"
-
-    use_block_sparsity = block_sparse_tensors is not None
-
-    if mask_mod is None:
-        if causal:
-            window_size_right = 0
-        local = window_size_left is not None or window_size_right is not None
-        if window_size_left is not None or window_size_right is not None:
-            if window_size_left is None and window_size_right == 0:
-                causal, local = True, False
-                window_size_right = None
-            else:
-                causal, local = False, True
-    else:
-        causal, local = False, False
-
-    current_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
-
-    if compute_capability == 9:  # TODO: tune block size according to hdim.
-        if (
-            head_dim == head_dim_v == 128
-            and not causal
-            and not local
-            and not use_block_sparsity
-        ):
-            n_block_size = 192
-
-    if compute_capability in [10, 11]:
-        if pack_gqa and (128 % qhead_per_kvhead != 0):
-            pack_gqa = False
-        # TODO: fix GQA + SplitKV + non-varlen
-        if pack_gqa and num_splits != 1 and cu_seqlens_q is None:
-            pack_gqa = False
-
-    if max_seqlen_q is None:
-        max_seqlen_q = seqlen_q if cu_seqlens_q is None else total_q
-    if max_seqlen_k is None:
-        max_seqlen_k = seqlen_k
-    seqlen_q_packgqa = max_seqlen_q * qhead_per_kvhead
-    if compute_capability == 10:
-        q_stage = 2 if seqlen_q_packgqa > m_block_size else 1
-    else:
-        q_stage = 1
-
-    if num_splits < 1:
-        m_block_size_effective = q_stage * m_block_size
-        seqlen_k_loaded = (
-            max_seqlen_k
-            if not local
-            else max(
-                0,
-                min(
-                    max_seqlen_k,
-                    window_size_right + window_size_left + 1 + m_block_size,
-                ),
-            )
-        )
-        num_n_blocks = (seqlen_k_loaded + n_block_size - 1) // n_block_size
-        num_m_blocks = (
-            seqlen_q_packgqa + m_block_size_effective - 1
-        ) // m_block_size_effective
-        total_mblocks = batch_size * num_head_kv * num_m_blocks
-        num_splits = num_splits_heuristic(
-            total_mblocks,
-            torch.cuda.get_device_properties(device).multi_processor_count,
-            num_n_blocks,
-            128,
-        )
-
-    is_split_kv = num_splits > 1
-    if is_split_kv:
-        out_partial = torch.empty(
-            num_splits,
-            *q_batch_seqlen_shape,
-            num_head,
-            head_dim_v,
-            dtype=torch.float32,
-            device=device,
-        )
-        lse_partial = torch.empty(
-            num_splits, *lse_shape, dtype=torch.float32, device=device
-        )
-
-    # hash score and mask mods for compile cache
-    score_mod_hash = utils.hash_callable(score_mod) if score_mod is not None else False
-    mask_mod_hash = utils.hash_callable(mask_mod) if mask_mod is not None else False
-
-    if softcap is not None:
-        assert score_mod is None, "softcap and score_mod cannot be used together"
-        score_mod = utils.create_softcap_scoremod(softcap)
-
-    is_varlen = (
-        cu_seqlens_q is not None
-        or cu_seqlens_k is not None
-        or seqused_q is not None
-        or seqused_k is not None
-    )
-
-    if mask_mod is not None:
-        if is_varlen:
-            raise NotImplementedError(
-                "mask_mod with aux_tensors is not yet supported for varlen sequences. This will be fixed in a future PR."
-            )
-
-    if use_block_sparsity:
-        if is_varlen:
-            raise NotImplementedError(
-                "Block sparsity is not yet supported for varlen sequences. This will be fixed in a future PR."
-            )
-        # NB: pack_gqa requires block sparse head dim == 1 (broadcasted)
-        if pack_gqa and block_sparse_tensors.mask_block_cnt.shape[1] != 1:
-            pack_gqa = False
-        if is_split_kv:
-            raise NotImplementedError(
-                "Block sparsity is not yet supported with SplitKV. TODO: partition sparse block lists per split."
-            )
-
-    compile_key = (
-        dtype,
-        head_dim,
-        head_dim_v,
-        qhead_per_kvhead,
-        causal,
-        score_mod_hash,
-        mask_mod_hash,
-        use_block_sparsity,
-        len(aux_tensors) if aux_tensors is not None else 0,
-        lse is None,
-        cu_seqlens_q is None,
-        cu_seqlens_k is None,
-        seqused_q is None,
-        seqused_k is None,
-        page_table is not None,
-        window_size_left is not None,
-        window_size_right is not None,
-        learnable_sink is not None,
-        m_block_size,
-        n_block_size,
-        q_stage,
-        num_threads,
-        is_split_kv,
-        pack_gqa,
-        compute_capability,
-        page_size not in [None, 128],  # paged KV non-TMA
-    )
-    if compile_key not in _flash_attn_fwd.compile_cache:
-        (
-            cu_seqlens_q_tensor,
-            cu_seqlens_k_tensor,
-            seqused_q_tensor,
-            seqused_k_tensor,
-            learnable_sink_tensor,
-        ) = [
-            to_cute_tensor(t, assumed_align=4, leading_dim=0) if t is not None else None
-            for t in (cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, learnable_sink)
-        ]
-        page_table_tensor = (
-            to_cute_tensor(page_table, assumed_align=4, leading_dim=1)
-            if page_table is not None
-            else None
-        )
-        q_tensor, k_tensor, v_tensor, o_tensor = [
-            to_cute_tensor(t)
-            for t in (q, k, v, out if not is_split_kv else out_partial)
-        ]
-        if is_split_kv:
-            lse_tensor = to_cute_tensor(lse_partial, assumed_align=4)
-        elif lse is not None:
-            lse_tensor = to_cute_tensor(lse, assumed_align=4)
-        else:
-            lse_tensor = None
-
-        sparse_tensors = None
-        if block_sparse_tensors is not None:
-            if seqlen_q is None:
-                raise ValueError(
-                    "Block sparsity requires fixed-length sequences (seqlen_q must be known)."
-                )
-            expected_count_shape, expected_index_shape = (
-                get_block_sparse_expected_shapes(
-                    batch_size,
-                    num_head,
-                    seqlen_q,
-                    seqlen_k,
-                    m_block_size,
-                    n_block_size,
-                    q_stage,
-                )
-            )
-            compile_time_normalized = normalize_block_sparse_tensors(
-                block_sparse_tensors,
-                expected_count_shape=expected_count_shape,
-                expected_index_shape=expected_index_shape,
-            )
-            sparse_tensors = to_cute_block_sparse_tensors(compile_time_normalized)
-
-        cute_aux_tensors = None
-        if aux_tensors is not None:
-            cute_aux_tensors = [
-                to_cute_tensor(buf, assumed_align=None, fully_dynamic=True)
-                for buf in aux_tensors
-            ]
-
-        if compute_capability == 9:
-            assert page_table is None, "paged KV not supported on SM 9.0"
-            assert not is_split_kv, "SplitKV not supported on SM 9.0"
-            # fa_fwd = FlashAttentionForwardSm80(
-            fa_fwd = FlashAttentionForwardSm90(
-                dtype,
-                head_dim,
-                head_dim_v,
-                qhead_per_kvhead,
-                is_causal=causal,
-                is_local=local,
-                pack_gqa=pack_gqa,
-                tile_m=m_block_size,
-                tile_n=n_block_size,
-                # num_stages=1,
-                num_stages=2,
-                num_threads=num_threads,
-                Q_in_regs=False,
-                intra_wg_overlap=True,
-                mma_pv_is_rs=True,
-                mask_mod=mask_mod,
-                score_mod=score_mod,
-                has_aux_tensors=aux_tensors is not None,
-            )
-        elif compute_capability in [10, 11]:
-            fa_fwd = FlashAttentionForwardSm100(
-                head_dim,
-                head_dim_v,
-                qhead_per_kvhead=qhead_per_kvhead,
-                is_causal=causal,
-                is_local=local,
-                is_split_kv=is_split_kv,
-                pack_gqa=pack_gqa,
-                m_block_size=m_block_size,
-                n_block_size=n_block_size,
-                q_stage=q_stage,
-                is_persistent=not causal
-                and not local
-                and cu_seqlens_q is None
-                and seqused_q is None
-                and not is_split_kv,
-                score_mod=score_mod,
-                mask_mod=mask_mod,
-                has_aux_tensors=aux_tensors is not None,
-                paged_kv_non_tma=page_size not in [None, 128],
-                is_varlen_q=cu_seqlens_q is not None or seqused_q is not None,
-            )
-        else:
-            raise ValueError(
-                f"Unsupported compute capability: {compute_capability}. Supported: 9.x, 10.x, 11.x"
-            )
-        # TODO: check @can_implement
-        _flash_attn_fwd.compile_cache[compile_key] = cute.compile(
-            fa_fwd,
-            q_tensor,
-            k_tensor,
-            v_tensor,
-            o_tensor,
-            lse_tensor,
-            softmax_scale,
-            current_stream,
-            cu_seqlens_q_tensor,
-            cu_seqlens_k_tensor,
-            seqused_q_tensor,
-            seqused_k_tensor,
-            page_table_tensor,
-            window_size_left,
-            window_size_right,
-            learnable_sink_tensor,
-            sparse_tensors,
-            cute_aux_tensors,
-            options="--enable-tvm-ffi",
-        )
-
-    # Expand block sparse tensors to match actual head count (may be broadcast from 1)
-    normalized_block_sparse_tensors = None
-    if block_sparse_tensors is not None:
-        expected_count_shape, expected_index_shape = get_block_sparse_expected_shapes(
-            batch_size,
-            num_head,
-            seqlen_q,
-            seqlen_k,
-            m_block_size,
-            n_block_size,
-            q_stage,
-        )
-        normalized_block_sparse_tensors = normalize_block_sparse_tensors(
-            block_sparse_tensors,
-            expected_count_shape=expected_count_shape,
-            expected_index_shape=expected_index_shape,
-        )
-    _flash_attn_fwd.compile_cache[compile_key](
-        q,
-        k,
-        v,
-        out if not is_split_kv else out_partial,
-        lse_partial if is_split_kv else lse,
-        softmax_scale,
-        current_stream,
-        cu_seqlens_q,
-        cu_seqlens_k,
-        seqused_q,
-        seqused_k,
-        page_table,
-        window_size_left,
-        window_size_right,
-        learnable_sink,
-        normalized_block_sparse_tensors,
-        aux_tensors,
-    )
-    if is_split_kv:
-        _flash_attn_fwd_combine(
-            out_partial,
-            lse_partial.transpose(-1, -2),
-            out,
-            lse.transpose(-1, -2) if lse is not None else None,
-            cu_seqlens_q,
-            seqused_q,
-        )
-    return out, lse
-
-
-_flash_attn_fwd.compile_cache = {}
-
-
-def _flash_attn_fwd_combine(
-    out_partial: torch.Tensor,
-    lse_partial: torch.Tensor,
-    out: torch.Tensor,
-    lse: Optional[torch.Tensor] = None,
-    cu_seqlens: Optional[torch.Tensor] = None,
-    seqused: Optional[torch.Tensor] = None,
-    num_splits_dynamic_ptr: Optional[torch.Tensor] = None,
-    semaphore_to_reset: Optional[torch.Tensor] = None,
-) -> None:
-    """Forward combine kernel for split attention computation.
-
-    Combines partial outputs and log-sum-exp values from multiple splits
-    of attention computation into final outputs.
-
-    Args:
-        out_partial: Partial outputs tensor (num_splits, batch, seqlen, nheads, headdim) or
-                                            (num_splits, total_q, nheads, headdim) if there's cu_seqlens
-        lse_partial: Partial LSE tensor (num_splits, batch, seqlen, nheads) or
-                                       (num_splits, total_q, nheads) if there's cu_seqlens
-        out: Output tensor (batch, seqlen, nheads, headdim) or (total_q, nheads, headdim) if there's cu_seqlens
-        lse: Output LSE tensor (batch, seqlen, nheads) or (total_q, nheads) if there's cu_seqlens.
-        cu_seqlens: Cumulative sequence lengths for variable length sequences
-        seqused: Used sequence lengths for each batch
-        num_splits_dynamic_ptr: Dynamic number of splits per batch
-        semaphore_to_reset: Semaphore for synchronization
-        k_block_size: Block size for head dimension
-
-    Returns:
-        None
-    """
-    # Input validation
-    assert out_partial.dim() in [4, 5], "out_partial must have 4 or 5 dimensions"
-    assert lse_partial.dim() in [3, 4], "lse_partial must have 3 or 4 dimensions"
-    assert out_partial.dtype in [
-        torch.float16,
-        torch.bfloat16,
-        torch.float32,
-    ], "out_partial must be fp16, bf16, or fp32"
-    assert lse_partial.dtype == torch.float32, "lse_partial must be fp32"
-    assert out_partial.is_cuda and lse_partial.is_cuda, "tensors must be on CUDA device"
-    assert (
-        out_partial.stride(-1) == 1
-    ), "out_partial must be contiguous in the last dimension"
-    assert (
-        lse_partial.stride(-2) == 1
-    ), "lse_partial must be contiguous in the seqlen dimension"
-    assert lse_partial.shape == out_partial.shape[:-1]
-
-    # Determine if this is variable length based on dimensions
-    is_varlen = out_partial.dim() == 4
-
-    # Validate output tensor shapes and types
-    assert out.shape == out_partial.shape[1:], "out shape mismatch"
-    if lse is not None:
-        assert lse.shape == lse_partial.shape[1:], "lse shape mismatch"
-        assert lse.dtype == torch.float32, "lse must be fp32"
-
-    # Validate optional tensors
-    for t, name in [
-        (cu_seqlens, "cu_seqlens"),
-        (seqused, "seqused"),
-        (num_splits_dynamic_ptr, "num_splits_dynamic_ptr"),
-    ]:
-        if t is not None:
-            assert t.dtype == torch.int32, f"{name} must be int32"
-            assert t.is_cuda, f"{name} must be on CUDA device"
-            assert t.is_contiguous(), f"{name} must be contiguous"
-
-    head_dim = out_partial.shape[-1]
-    num_splits = out_partial.shape[0]
-    assert num_splits <= 256
-    # If hdim is 96 or 192, it's faster to round them to 128 or 256 respectively
-    # so that kBlockM is smaller and we have more parallelism.
-    k_block_size = 64 if head_dim <= 64 else 128
-    # We want kBlockM to be as small as possible to maximize parallelism.
-    # E.g., if hdim is 64, we want kBlockM to be 16 so that we can use 256 threads, each reading 4 elements (floats).
-    m_block_size = (
-        8 if k_block_size % 128 == 0 else (16 if k_block_size % 64 == 0 else 32)
-    )
-    log_max_splits = max(math.ceil(math.log2(num_splits)), 4)
-    if m_block_size == 8:
-        # If kBlockM == 8 then the minimum number of splits is 32.
-        # TODO: we can deal w this by using 128 threads instead
-        log_max_splits = max(log_max_splits, 5)
-
-    current_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
-
-    # Create combine kernel configuration
-    dtype = torch2cute_dtype_map[out.dtype]
-    dtype_partial = torch2cute_dtype_map[out_partial.dtype]
-
-    compile_key = (
-        dtype,
-        dtype_partial,
-        head_dim,
-        m_block_size,
-        k_block_size,
-        log_max_splits,
-        cu_seqlens is not None,
-        seqused is not None,
-        lse is not None,
-    )
-
-    if compile_key not in _flash_attn_fwd_combine.compile_cache:
-        out_partial_tensor = to_cute_tensor(
-            out_partial, leading_dim=4 if not is_varlen else 3
-        )
-        lse_partial_tensor = to_cute_tensor(
-            lse_partial, assumed_align=4, leading_dim=lse_partial.ndim - 2
-        )
-        out_tensor = to_cute_tensor(out, leading_dim=3 if not is_varlen else 2)
-        lse_tensor = (
-            to_cute_tensor(lse, assumed_align=4, leading_dim=lse.ndim - 2)
-            if lse is not None
-            else None
-        )
-
-        optional_tensors = [
-            to_cute_tensor(t, assumed_align=4, leading_dim=0) if t is not None else None
-            for t in (cu_seqlens, seqused, num_splits_dynamic_ptr, semaphore_to_reset)
-        ]
-        (
-            cu_seqlens_tensor,
-            seqused_tensor,
-            num_splits_dynamic_tensor,
-            semaphore_tensor,
-        ) = optional_tensors
-        fa_combine = FlashAttentionForwardCombine(
-            dtype=dtype,
-            dtype_partial=dtype_partial,
-            head_dim=head_dim,
-            m_block_size=m_block_size,
-            k_block_size=k_block_size,
-            log_max_splits=log_max_splits,
-        )
-
-        # Check if implementation is supported
-        if not fa_combine.can_implement(
-            dtype,
-            dtype_partial,
-            head_dim,
-            m_block_size,
-            k_block_size,
-            log_max_splits,
-            num_threads=256,
-        ):
-            raise RuntimeError(
-                "FlashAttention combine kernel cannot be implemented with given parameters"
-            )
-
-        _flash_attn_fwd_combine.compile_cache[compile_key] = cute.compile(
-            fa_combine,
-            out_partial_tensor,
-            lse_partial_tensor,
-            out_tensor,
-            lse_tensor,
-            cu_seqlens_tensor,
-            seqused_tensor,
-            num_splits_dynamic_tensor,
-            semaphore_tensor,
-            current_stream,
-            options="--enable-tvm-ffi",
-        )
-    _flash_attn_fwd_combine.compile_cache[compile_key](
-        out_partial,
-        lse_partial,
-        out,
-        lse,
-        cu_seqlens,
-        seqused,
-        num_splits_dynamic_ptr,
-        semaphore_to_reset,
-        current_stream,
-    )
-
-
-_flash_attn_fwd_combine.compile_cache = {}
-
-
-def warmup_flash_attn(f):
-    """
-    Decorator for flash_attn_varlen_func:
-    - On first call, run several warmup passes with different flag combinations:
-        * return_softmax_lse in {False, True}
-        * global noncausal (window_size=(None,None))
-        * causal (window_size=(None,0))
-        * local sliding window (window_size=(64,64))
-        * optionally pack_gqa=True if qheads > kvheads and allowed
-    - No score_mod / softcap (not supported for varlen yet)
-    - Executes sequentially to minimize peak GPU mem
-    - Does not modify user tensors (clones)
-    """
-    disable_warmup = os.getenv("SGLANG_DISABLE_FA4_WARMUP", "").lower() in (
-        "1",
-        "true",
-        "yes",
-        "on",
-    )
-    if disable_warmup:
-        return f
-
-    done = False
-
-    def _clone_args(args, kwargs):
-        """Clone tensor arguments to avoid sharing storage; deepcopy for others."""
-
-        def maybe_clone(x):
-            if isinstance(x, torch.Tensor):
-                return x.detach().clone()  # detach to avoid autograd edges
-            return copy.deepcopy(x)
-
-        return tuple(maybe_clone(a) for a in args), {
-            k: maybe_clone(v) for k, v in kwargs.items()
-        }
-
-    def _infer_heads(args, kwargs):
-        """Infer q and kv head counts from arguments."""
-        # Expect signature: (q, k, v, cu_seqlens_q, cu_seqlens_k, ...)
-        q = args[0] if len(args) > 0 else kwargs.get("q")
-        k = args[1] if len(args) > 1 else kwargs.get("k")
-        try:
-            qh = int(q.shape[-2])
-            kvh = int(k.shape[-2])
-            return qh, kvh
-        except Exception:
-            return None, None
-
-    def _run_warmups(args, kwargs):
-        """Run warmup calls sequentially and release memory after each."""
-        base_args, base_kwargs = _clone_args(args, kwargs)
-
-        qh, kvh = _infer_heads(base_args, base_kwargs)
-        can_pack_gqa = (
-            qh is not None and kvh is not None and qh % kvh == 0 and qh // kvh > 1
-        )
-        has_page_table = (
-            "page_table" in base_kwargs and base_kwargs["page_table"] is not None
-        )
-
-        # Window presets covering global, causal, and local
-        window_presets = [
-            (None, None),  # global noncausal
-            (None, 0),  # causal
-            (64, 64),  # local sliding window
-        ]
-
-        lse_flags = [False, True]
-
-        # Base combo list
-        combos = []
-        for ws in window_presets:
-            for return_lse_flag in lse_flags:
-                combos.append(dict(window_size=ws, return_softmax_lse=return_lse_flag))
-
-        # Optionally add a pack_gqa=True variant (FA4 may disable it internally for some varlen shapes/SMs)
-        if can_pack_gqa:
-            for ws in window_presets:
-                combos.append(
-                    dict(window_size=ws, return_softmax_lse=False, pack_gqa=True)
-                )
-
-        # If page_table is present, warm one combo with it (page_table in compile key for SM100)
-        if has_page_table:
-            combos.append(dict(window_size=(None, None), return_softmax_lse=False))
-
-        # Run sequentially
-        for combo in combos:
-            wa, wk = _clone_args(base_args, base_kwargs)
-            # Keep user-provided softcap/score_mod OUT (varlen+score_mod unsupported)
-            wk.pop("score_mod", None)
-            if "softcap" in wk and wk["softcap"]:
-                wk["softcap"] = 0.0
-            # Apply combo
-            wk.update(combo)
-            with torch.cuda.stream(torch.cuda.current_stream()):
-                try:
-                    f(*wa, **wk)
-                except Exception as e:
-                    # Some combos can be invalid for specific head dims / arch. Ignore and continue.
-                    logger.debug("Warmup combo skipped: %s", e)
-            del wa, wk
-            torch.cuda.empty_cache()
-            gc.collect()
-
-    def wrapper(*args, **kwargs):
-        nonlocal done
-        if not done:
-            logger.info(
-                "Running FA4 warmup (global/causal/local, LSE on/off, optional GQA pack)..."
-            )
-            _run_warmups(args, kwargs)
-            done = True
-        return f(*args, **kwargs)
-
-    return wrapper
-
-
-@warmup_flash_attn
-def flash_attn_varlen_func(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    cu_seqlens_q: Optional[torch.Tensor] = None,
-    cu_seqlens_k: Optional[torch.Tensor] = None,
-    seqused_q: Optional[torch.Tensor] = None,
-    seqused_k: Optional[torch.Tensor] = None,
-    page_table: Optional[torch.Tensor] = None,
-    softmax_scale: Optional[float] = None,
-    causal: bool = False,
-    window_size: Tuple[Optional[int], Optional[int]] = (None, None),
-    learnable_sink: Optional[torch.Tensor] = None,
-    softcap: float = 0.0,
-    num_splits: int = 1,
-    pack_gqa: Optional[bool] = None,
-    return_softmax_lse: Optional[bool] = False,
-    score_mod: Optional[Callable] = None,
-    aux_tensors: Optional[list] = None,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    out, lse = _flash_attn_fwd(
-        q,
-        k,
-        v,
-        cu_seqlens_q,
-        cu_seqlens_k,
-        seqused_q,
-        seqused_k,
-        page_table=page_table,
-        softmax_scale=softmax_scale,
-        causal=causal,
-        window_size_left=window_size[0],
-        window_size_right=window_size[1],
-        learnable_sink=learnable_sink,
-        softcap=softcap,
-        num_splits=num_splits,
-        pack_gqa=pack_gqa,
-        return_lse=return_softmax_lse,
-        score_mod=score_mod,
-        aux_tensors=aux_tensors,
-    )
-
-    return (out, lse) if return_softmax_lse else out

sgl-kernel/python/sgl_kernel/attention.py

@@ -3,25 +3,6 @@ from typing import Optional, Tuple
 import torch
 
 
-def merge_state(
-    v_a: torch.Tensor,
-    s_a: torch.Tensor,
-    v_b: torch.Tensor,
-    s_b: torch.Tensor,
-    v_merged: Optional[torch.Tensor] = None,
-    s_merged: Optional[torch.Tensor] = None,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    s_a = s_a.to(torch.float32)
-    s_b = s_b.to(torch.float32)
-    # Avoid creating new tensors if they are already provided
-    if v_merged is None:
-        v_merged = torch.empty_like(v_a)
-    if s_merged is None:
-        s_merged = torch.empty_like(s_a)
-    torch.ops.sgl_kernel.merge_state.default(v_a, s_a, v_b, s_b, v_merged, s_merged)
-    return v_merged, s_merged
-
-
 def merge_state_v2(
     v_a: torch.Tensor,
     s_a: torch.Tensor,

sgl-kernel/python/sgl_kernel/elementwise.py

@@ -344,22 +344,6 @@ def rotary_embedding(
     )
 
 
-def downcast_fp8(
-    k: torch.Tensor,
-    v: torch.Tensor,
-    k_out: torch.Tensor,
-    v_out: torch.Tensor,
-    k_scale: torch.Tensor,
-    v_scale: torch.Tensor,
-    loc: torch.Tensor,
-    mult: int = 1,
-    offset: int = 0,
-) -> None:
-    torch.ops.sgl_kernel.downcast_fp8(
-        k, v, k_out, v_out, k_scale, v_scale, loc, mult, offset
-    )
-
-
 def copy_to_gpu_no_ce(input: torch.Tensor, output: torch.Tensor):
     torch.ops.sgl_kernel.copy_to_gpu_no_ce(input, output)
 

sgl-kernel/python/sgl_kernel/sampling.py

@@ -113,76 +113,3 @@ def top_p_renorm_probs(
 
 
 top_p_renorm_prob = top_p_renorm_probs
-
-
-def _top_k_mask_logits_internal(
-    logits: torch.Tensor,
-    maybe_top_k_arr: Optional[torch.Tensor],
-    top_k_val: int,
-) -> torch.Tensor:
-    logits = logits.float()
-    maybe_top_k_arr = maybe_top_k_arr.int() if maybe_top_k_arr is not None else None
-    mask_logits = torch.empty_like(logits)
-    torch.ops.sgl_kernel.top_k_mask_logits.default(
-        logits, mask_logits, maybe_top_k_arr, top_k_val
-    )
-    return mask_logits
-
-
-def top_k_mask_logits(
-    logits: torch.Tensor,
-    top_k: Union[torch.Tensor, int],
-) -> torch.Tensor:
-    r"""Adapt from https://github.com/flashinfer-ai/flashinfer/flashinfer/sampling.py
-    Fused GPU kernel for masking logits by top-k thresholding.
-
-    Parameters
-    ----------
-    logits: torch.Tensor
-        Logits before softmax, shape ``(batch_size, num_classes)``.
-    top_k: Union[torch.Tensor, int]
-        Either a scalar or a tensor of shape ``(batch_size,)``, representing the top-k threshold for for
-        for masking logits, should be in ``(0, num_classes)``.
-        If a scalar, the same threshold is used for all requests.
-        If a tensor, each request has its own threshold.
-        We keep the top-k logits, set the rest to negative infinity.
-
-    Returns
-    -------
-    masked_logits: torch.Tensor
-        Masked logits, shape ``(batch_size, num_classes)``.
-
-    Examples
-    --------
-
-    >>> import torch
-    >>> import flashinfer
-    >>> torch.manual_seed(42)
-    >>> batch_size = 4
-    >>> vocab_size = 5
-    >>> top_k = 3
-    >>> logits = torch.randn(batch_size, vocab_size).to(0)
-    >>> logits
-    tensor([[ 1.9269,  1.4873,  0.9007, -2.1055, -0.7581],
-            [ 1.0783,  0.8008,  1.6806,  0.3559, -0.6866],
-            [-0.4934,  0.2415, -0.2316,  0.0418, -0.2516],
-            [ 0.8599, -0.3097, -0.3957,  0.8034, -0.6216]], device='cuda:0')
-    >>> masked_logits = flashinfer.sampling.top_k_mask_logits(logits, top_k)
-    >>> masked_logits
-    tensor([[ 1.9269,  1.4873,  0.9007,    -inf,    -inf],
-            [ 1.0783,  0.8008,  1.6806,    -inf,    -inf],
-            [   -inf,  0.2415, -0.2316,  0.0418,    -inf],
-            [ 0.8599, -0.3097,    -inf,  0.8034,    -inf]], device='cuda:0')
-
-    Note
-    ----
-    The combination of ``top_k_mask_logits`` and ``softmax`` should be equivalent to ``top_k_renorm_probs``.
-
-    See Also
-    --------
-    top_k_renorm_probs
-    """
-    if logits.device.type == "musa" or not _has_flashinfer:
-        return _top_k_mask_logits_internal(logits, *_to_tensor_scalar_tuple(top_k))
-    else:
-        return _flashinfer_sampling.top_k_mask_logits(logits, top_k)

sgl-kernel/tests/test_hadamard.py

@@ -1,86 +0,0 @@
-import math
-import sys
-
-import pytest
-import torch
-import torch.nn.functional as F
-from einops import rearrange, repeat
-from scipy.linalg import hadamard
-
-try:
-    from sgl_kernel import hadamard_transform
-except Exception:
-    pytest.skip(
-        "sgl-kernel hadamard interface was removed (migrated to jit_kernel)",
-        allow_module_level=True,
-    )
-
-
-def hadamard_transform_ref(x, scale=1.0):
-    """
-    x: (..., dim)
-    out: (..., dim)
-    """
-    if hadamard is None:
-        raise ImportError("Please install scipy")
-    x_shape = x.shape
-    dim = x.shape[-1]
-    x = x.reshape(-1, dim)
-    log_dim = math.ceil(math.log2(dim))
-    dim_padded = 2**log_dim
-    if dim != dim_padded:
-        x = F.pad(x, (0, dim_padded - dim))
-    out = F.linear(
-        x,
-        torch.tensor(hadamard(dim_padded, dtype=float), dtype=x.dtype, device=x.device),
-    )
-    out = out * scale
-    return out[..., :dim].reshape(*x_shape)
-
-
-@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
-@pytest.mark.parametrize(
-    "dim",
-    [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 137, 1024, 2048, 4096, 8192, 16384, 32768],
-)
-def test_fast_hadamard_transform(dim, dtype):
-    device = "cuda"
-
-    if dtype == torch.float32:
-        rtol, atol = 3e-4, 3e-3
-    elif dtype == torch.bfloat16:
-        rtol, atol = 1e-2, 5e-2
-    else:  # float16
-        rtol, atol = 3e-3, 5e-3
-
-    torch.random.manual_seed(0)
-    batch_size = 15
-
-    x = torch.randn(batch_size, dim, device=device, dtype=dtype)
-    x_ref = x.detach().clone().to(torch.float32)
-    x_pt = x.detach().clone()
-
-    scale = 1 / math.sqrt(dim)
-
-    out = hadamard_transform(x, scale=scale)
-    out_ref = hadamard_transform_ref(x_ref, scale=scale)
-    out_pt = hadamard_transform_ref(x_pt, scale=scale)
-
-    torch.testing.assert_close(
-        out_pt.float(),
-        out_ref,
-        rtol=rtol,
-        atol=atol,
-        msg="Reference implementations mismatch",
-    )
-    torch.testing.assert_close(
-        out.float(),
-        out_ref,
-        rtol=rtol,
-        atol=atol,
-        msg="fast_hadamard_transform output mismatch",
-    )
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__]))

sgl-kernel/tests/test_merge_state.py

@@ -1,143 +0,0 @@
-# Adapted from https://github.com/flashinfer-ai/flashinfer/blob/55576c626421b5ee7e7ebe74afd26465c8ae863f/flashinfer/triton/kernels/cascade.py
-
-import sys
-from typing import List
-
-import pytest
-import torch
-import triton
-import triton.language as tl
-from sgl_kernel import merge_state
-
-
-def check_input(x: torch.Tensor):
-    assert x.is_cuda, f"{str(x)} must be a CUDA Tensor"
-    assert x.is_contiguous(), f"{str(x)} must be contiguous"
-
-
-def check_dim(d, x: torch.Tensor):
-    assert x.dim() == d, f"{str(x)} must be a {d}D tensor"
-
-
-def check_shape(a: torch.Tensor, b: torch.Tensor):
-    assert a.dim() == b.dim(), "tensors should have same dim"
-    for i in range(a.dim()):
-        assert a.size(i) == b.size(
-            i
-        ), f"tensors shape mismatch, {a.size()} and {b.size()}"
-
-
-def check_device(tensors: List[torch.Tensor]):
-    device = tensors[0].device
-    for t in tensors:
-        assert (
-            t.device == device
-        ), f"All tensors should be on the same device, but got {device} and {t.device}"
-
-
-@triton.jit
-def state_merge(o, m, d, other_o, other_m, other_d):
-    m_max = tl.maximum(m, other_m)
-    d = d * tl.exp2(m - m_max) + other_d * tl.exp2(other_m - m_max)
-    o = o * tl.exp2(m - m_max) + other_o * tl.exp2(other_m - m_max)
-    return o, m_max, d
-
-
-@triton.jit
-def state_normalize(o, m, d):
-    o = o / d
-    return o, m, d
-
-
-@triton.jit
-def state_get_lse(o, m, d):
-    return m + tl.log2(d)
-
-
-@triton.jit
-def merge_state_kernel(
-    v_a_ptr,
-    s_a_ptr,
-    v_b_ptr,
-    s_b_ptr,
-    v_merged_ptr,
-    s_merged_ptr,
-    num_heads,
-    head_dim,
-    bdx: tl.constexpr,
-    bdy: tl.constexpr,
-):
-    pos = tl.program_id(axis=0)
-    for tx in tl.range(bdx):
-        for head_idx in tl.range(bdy):
-            s_a_val = tl.load(s_a_ptr + pos * num_heads + head_idx)
-            s_b_val = tl.load(s_b_ptr + pos * num_heads + head_idx)
-
-            offsets = (pos * num_heads + head_idx) * head_dim + tx
-            v_a = tl.load(v_a_ptr + offsets)
-            v_b = tl.load(v_b_ptr + offsets)
-
-            v_merged, s_max, d = state_merge(
-                o=v_a, m=s_a_val, d=1, other_o=v_b, other_m=s_b_val, other_d=1
-            )
-            v_merged, s_max, d = state_normalize(v_merged, s_max, d)
-            v_merged_offset = (pos * num_heads + head_idx) * head_dim + tx
-            tl.store(v_merged_ptr + v_merged_offset, v_merged)
-
-            if s_merged_ptr:
-                tl.store(
-                    s_merged_ptr + pos * num_heads + head_idx,
-                    tl.log2(d) + s_max,
-                )
-
-
-def merge_state_triton(
-    v_a: torch.Tensor, s_a: torch.Tensor, v_b: torch.Tensor, s_b: torch.Tensor
-):
-    check_input(v_a)
-    check_input(s_a)
-    check_input(v_b)
-    check_input(s_b)
-    check_device([v_a, s_a, v_b, s_b])
-    check_dim(3, v_a)
-    check_dim(2, s_a)
-    check_dim(3, v_b)
-    check_dim(2, s_b)
-    check_shape(v_a, v_b)
-    check_shape(s_a, s_b)
-    assert v_a.size(0) == s_a.size(0)
-    assert v_a.size(1) == s_b.size(1)
-    s_a = s_a.to(torch.float32)
-    s_b = s_b.to(torch.float32)
-    seq_len = v_a.size(0)
-    num_heads = v_a.size(1)
-    head_dim = v_a.size(2)
-    v_merged = torch.empty_like(v_a).to(s_a.device)
-    s_merged = torch.empty((seq_len, num_heads)).to(s_a.device)
-    bdx = head_dim
-    bdy = num_heads
-
-    merge_state_kernel[lambda meta: (seq_len,)](
-        v_a, s_a, v_b, s_b, v_merged, s_merged, num_heads, head_dim, bdx=bdx, bdy=bdy
-    )
-
-    return v_merged, s_merged
-
-
-@pytest.mark.parametrize("seq_len", [2048])
-@pytest.mark.parametrize("num_heads", [32])
-@pytest.mark.parametrize("head_dim", [128])
-def test_merge_state(seq_len, num_heads, head_dim):
-    va = torch.randn(seq_len, num_heads, head_dim).half().to("cuda:0")
-    sa = torch.randn(seq_len, num_heads, dtype=torch.float32).to("cuda:0")
-    vb = torch.randn(seq_len, num_heads, head_dim).half().to("cuda:0")
-    sb = torch.randn(seq_len, num_heads, dtype=torch.float32).to("cuda:0")
-    v_merged, s_merged = merge_state_triton(va, sa, vb, sb)
-    v_merged_std, s_merged_std = merge_state(va, sa, vb, sb)
-
-    assert torch.allclose(v_merged, v_merged_std, atol=1e-2)
-    assert torch.allclose(s_merged, s_merged_std, atol=1e-2)
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__]))

sgl-kernel/tests/test_merge_state_v2.py

@@ -5,7 +5,7 @@ import pytest
 import torch
 import triton
 import triton.language as tl
-from sgl_kernel import merge_state, merge_state_v2
+from sgl_kernel import merge_state_v2
 
 
 @triton.jit
@@ -146,11 +146,9 @@ def generate_markdown_table():
     global all_case_info
     table_header = (
         "| tokens | heads | headsize | dtype "
-        "| device | torch | triton | v1 | v2 | speedup(vs triton) | speedup(vs v1)|"
+        "| device | torch | triton | v2 | speedup(vs triton) |"
-    )
-    table_separator = (
-        "| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |"
     )
+    table_separator = "| --- | --- | --- | --- | --- | --- | --- | --- | --- |"
 
     def shortly_dtype(dtype: torch.dtype) -> str:
         return str(dtype).removeprefix("torch.")
@@ -169,21 +167,17 @@ def generate_markdown_table():
             device,
             time_torch,
             time_triton,
-            time_v1,
             time_v2,
         ) = info
         dtype = shortly_dtype(dtype)
         device = shortly_device(device)
         improved_triton = time_triton / time_v2
-        improved_v1 = time_v1 / time_v2
         print(
             f"| {num_tokens} | {num_heads} | {head_size} "
             f"| {dtype} | {device} | {time_torch:.4f}ms "
             f"| {time_triton:.4f}ms "
-            f"| {time_v1:.4f}ms "
             f"| {time_v2:.4f}ms "
-            f"| {improved_triton:.4f}x "
+            f"| {improved_triton:.4f}x |"
-            f"| {improved_v1:.4f}x |"
         )
 
 
@@ -259,11 +253,6 @@ def test_merge_attn_states(
             prefix_lse_ = prefix_lse
             suffix_lse_ = suffix_lse
 
-        if fn_type == "cuda_v1":
-            # merge_state v1 kernel not support float32
-            if output_dtype not in (torch.half, torch.bfloat16):
-                return 0, output_fn, output_lse_fn
-
         total_time = 0
         start = torch.cuda.Event(enable_timing=True)
         end = torch.cuda.Event(enable_timing=True)
@@ -316,29 +305,21 @@ def test_merge_attn_states(
         fn_type="triton",
     )
 
-    # 2. Run the merge_state V1 kernel
+    # 2. Run the merge_state V2 kernel
-    output_v1 = output.clone()
-    output_lse_v1 = output_lse.clone()
-    time_v1, output_v1, output_lse_v1 = perf_kernel_fn(
-        output_v1, output_lse_v1, merge_state, fn_type="cuda_v1"
-    )
-
-    # 3. Run the merge_state V2 kernel
     output_v2 = output.clone()
     output_lse_v2 = output_lse.clone()
     time_v2, output_v2, output_lse_v2 = perf_kernel_fn(
         output_v2, output_lse_v2, merge_state_v2, fn_type="cuda_v2"
     )
 
-    # 4. Performance compare
+    # 3. Performance compare
     improved = time_triton / time_v2
     print(f"  Torch time: {time_torch:.6f}ms")
     print(f" Triton time: {time_triton:.6f}ms")
-    print(f"CUDA v1 time: {time_v1:.6f}ms")
     print(f"CUDA v2 time: {time_v2:.6f}ms, Performance: {improved:.5f}x")
     print("-" * 100)
 
-    # 5. Correctness compare
+    # 4. Correctness compare
     # Liger Kernel: Efficient Triton Kernels for LLM Training
     # https://arxiv.org/pdf/2410.10989, 3.3 Correctness
     # use rtol = 1e-2 for bfloat16.
@@ -387,7 +368,6 @@ def test_merge_attn_states(
             device,
             time_torch,
             time_triton,
-            time_v1,
             time_v2,
         )
     )