diff --git a/.github/labeler.yml b/.github/labeler.yml index 2f1a92915..0e511e4fa 100644 --- a/.github/labeler.yml +++ b/.github/labeler.yml @@ -92,7 +92,6 @@ blackwell: - changed-files: - any-glob-to-any-file: - '**/*nvfp4*' - - 'python/sglang/kernels/aot/csrc/attention/cutlass_sm100_mla/**/*' - 'python/sglang/srt/layers/attention/trtllm_mla_backend.py' - 'python/sglang/srt/layers/attention/trtllm_mha_backend.py' diff --git a/docs/cookbook/autoregressive/DeepSeek/DeepSeek-V3.mdx b/docs/cookbook/autoregressive/DeepSeek/DeepSeek-V3.mdx index 61fc53f6d..27c65ca95 100644 --- a/docs/cookbook/autoregressive/DeepSeek/DeepSeek-V3.mdx +++ b/docs/cookbook/autoregressive/DeepSeek/DeepSeek-V3.mdx @@ -404,7 +404,7 @@ The spec-v2 overlap scheduler is enabled by default. It improves performance by DeepSeek V3 uses [Multi-head Latent Attention (MLA)](https://arxiv.org/pdf/2405.04434), an attention mechanism that improves inference efficiency. SGLang implements several optimizations: - **Weight Absorption:** Reorders matrix multiplications to improve decoding phase efficiency. -- **MLA Attention Backends:** FA3, Flashinfer, FlashMLA, CutlassMLA, TRTLLM MLA (Blackwell), and Triton. FA3 is the default. +- **MLA Attention Backends:** FA3, Flashinfer, FlashMLA, TRTLLM MLA (Blackwell), and Triton. FA3 is the default. - **FP8 Quantization:** W8A8 FP8 and KV Cache FP8, with BMM operators for weight-absorbed MLA in FP8. - **CUDA Graph & Torch.compile:** Both MLA and MoE support CUDA Graph and Torch.compile for reduced decoding latency. - **Chunked Prefix Cache:** Increases throughput for long-sequence chunked prefill (FlashAttention3 backend only). diff --git a/docs/docs/advanced_features/attention_backend.mdx b/docs/docs/advanced_features/attention_backend.mdx index ef7fcb6f3..04e70b6de 100644 --- a/docs/docs/advanced_features/attention_backend.mdx +++ b/docs/docs/advanced_features/attention_backend.mdx @@ -664,15 +664,6 @@ python3 -m sglang.launch_server \ --trust-remote-code ``` -- Cutlass MLA -```bash Command -python3 -m sglang.launch_server \ - --tp 8 \ - --model deepseek-ai/DeepSeek-R1 \ - --attention-backend cutlass_mla \ - --trust-remote-code -``` - - Ascend ```bash Command python3 -m sglang.launch_server \ @@ -701,13 +692,6 @@ python3 -m sglang.launch_server \ --attention-backend flex_attention ``` -- Dual Chunk FlashAttention -```bash Command -python3 -m sglang.launch_server \ - --model Qwen/Qwen2.5-14B-Instruct-1M \ - --attention-backend dual_chunk_flash_attn -``` - - Torch Native ```bash Command python3 -m sglang.launch_server \ diff --git a/docs/docs/advanced_features/quantization.mdx b/docs/docs/advanced_features/quantization.mdx index ccbf5526d..befcb4d93 100644 --- a/docs/docs/advanced_features/quantization.mdx +++ b/docs/docs/advanced_features/quantization.mdx @@ -89,14 +89,14 @@ The following table summarizes quantization method support across NVIDIA and AMD Yes Yes Yes - Uses Triton dequantize on AMD (vs. optimized CUDA kernels on NVIDIA). Uses CANN kernels on Ascend + Uses a JIT-compiled CUDA kernel on NVIDIA, Triton dequantize on AMD. Uses CANN kernels on Ascend gptq + No + No Yes - Yes - Yes - Uses Triton or vLLM kernels on AMD. Uses CANN kernels on Ascend + Removed on NVIDIA and AMD GPUs — use gptq_marlin instead. Uses CANN kernels on Ascend. Still supported on Intel CPUs with AMX compressed-tensors diff --git a/docs/docs/advanced_features/server_arguments.mdx b/docs/docs/advanced_features/server_arguments.mdx index f6080eeee..a10536533 100644 --- a/docs/docs/advanced_features/server_arguments.mdx +++ b/docs/docs/advanced_features/server_arguments.mdx @@ -1418,19 +1418,19 @@ Please consult the documentation below and [server_args.py](https://github.com/s `--attention-backend` Choose the kernels for attention layers. `None` - triton, torch_native, flex_attention, dsa, nsa, dsv4, compressed, cutlass_mla, fa3, fa4, flashinfer, flashmla, trtllm_mla, cutedsl_mla, tokenspeed_mla, trtllm_mha, dual_chunk_flash_attn, aiter, wave, intel_amx, ascend, intel_xpu + triton, torch_native, flex_attention, dsa, nsa, dsv4, compressed, fa3, fa4, flashinfer, flashmla, trtllm_mla, cutedsl_mla, tokenspeed_mla, trtllm_mha, aiter, wave, intel_amx, ascend, intel_xpu `--prefill-attention-backend` Choose the kernels for prefill attention layers (have priority over --attention-backend). `None` - triton, torch_native, flex_attention, dsa, nsa, dsv4, compressed, cutlass_mla, fa3, fa4, flashinfer, flashmla, trtllm_mla, cutedsl_mla, tokenspeed_mla, trtllm_mha, dual_chunk_flash_attn, aiter, wave, intel_amx, ascend, intel_xpu + triton, torch_native, flex_attention, dsa, nsa, dsv4, compressed, fa3, fa4, flashinfer, flashmla, trtllm_mla, cutedsl_mla, tokenspeed_mla, trtllm_mha, aiter, wave, intel_amx, ascend, intel_xpu `--decode-attention-backend` Choose the kernels for decode attention layers (have priority over --attention-backend). `None` - triton, torch_native, flex_attention, dsa, nsa, dsv4, compressed, cutlass_mla, fa3, fa4, flashinfer, flashmla, trtllm_mla, cutedsl_mla, tokenspeed_mla, trtllm_mha, dual_chunk_flash_attn, aiter, wave, intel_amx, ascend, intel_xpu + triton, torch_native, flex_attention, dsa, nsa, dsv4, compressed, fa3, fa4, flashinfer, flashmla, trtllm_mla, cutedsl_mla, tokenspeed_mla, trtllm_mha, aiter, wave, intel_amx, ascend, intel_xpu `--sampling-backend` diff --git a/examples/runtime/engine/offline_batch_inference_qwen_1m.py b/examples/runtime/engine/offline_batch_inference_qwen_1m.py deleted file mode 100644 index 5505bf7b4..000000000 --- a/examples/runtime/engine/offline_batch_inference_qwen_1m.py +++ /dev/null @@ -1,71 +0,0 @@ -""" -Usage: -python3 offline_batch_inference.py -""" - -from urllib.request import urlopen - -import sglang as sgl - - -def load_prompt() -> str: - # Test cases with various lengths can be found at: - # - # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/64k.txt - # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/200k.txt - # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/600k.txt - # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/1m.txt - - with urlopen( - "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/64k.txt", - timeout=5, - ) as response: - prompt = response.read().decode("utf-8") - return prompt - - -# Processing the prompt. -def process_requests(llm: sgl.Engine, prompts: list[str]) -> None: - # Create a sampling params object. - sampling_params = { - "temperature": 0.7, - "top_p": 0.8, - "top_k": 20, - "repetition_penalty": 1.05, - "max_new_tokens": 256, - } - # Generate texts from the prompts. - outputs = llm.generate(prompts, sampling_params) - # Print the outputs. - for output in outputs: - prompt_token_ids = output["meta_info"]["prompt_tokens"] - generated_text = output["text"] - print(f"Prompt length: {prompt_token_ids}, Generated text: {generated_text!r}") - - -# Create an LLM. -def initialize_engine() -> sgl.Engine: - llm = sgl.Engine( - model_path="Qwen/Qwen2.5-7B-Instruct-1M", - context_length=1048576, - page_size=256, - attention_backend="dual_chunk_flash_attn", - tp_size=4, - disable_radix_cache=True, - enable_mixed_chunk=False, - enable_torch_compile=False, - chunked_prefill_size=131072, - mem_fraction_static=0.6, - log_level="DEBUG", - ) - return llm - - -def main(): - llm = initialize_engine() - prompt = load_prompt() - process_requests(llm, [prompt]) - - -if __name__ == "__main__": - main() diff --git a/python/sglang/kernels/aot/CMakeLists.txt b/python/sglang/kernels/aot/CMakeLists.txt index 26943a769..4db9b4b15 100644 --- a/python/sglang/kernels/aot/CMakeLists.txt +++ b/python/sglang/kernels/aot/CMakeLists.txt @@ -256,9 +256,7 @@ endif() # NOTE: Please sort the filenames alphabetically set(SOURCES "csrc/allreduce/custom_all_reduce.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/concat_mla.cu" @@ -271,13 +269,11 @@ set(SOURCES "csrc/expert_specialization/es_sm100_mxfp8_blockscaled.cu" "csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cu" - "csrc/gemm/awq_kernel.cu" "csrc/gemm/fp8_gemm_kernel.cu" "csrc/gemm/int8_gemm_kernel.cu" "csrc/gemm/per_token_group_quant_8bit.cu" "csrc/gemm/per_token_group_quant_8bit_v2.cu" "csrc/gemm/per_token_quant_fp8.cu" - "csrc/gemm/gptq/gptq_kernel.cu" "csrc/grammar/apply_token_bitmask_inplace_cuda.cu" "csrc/infllm_v2/max_pooling.cu" @@ -306,12 +302,6 @@ set(SOURCES "${repo-flashinfer_SOURCE_DIR}/csrc/norm.cu" "${repo-flashinfer_SOURCE_DIR}/csrc/renorm.cu" - - "${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_sparse_hdim128_bf16_causal_sm80.cu" - "${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_sparse_hdim128_bf16_sm80.cu" - "${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_sparse_hdim128_fp16_causal_sm80.cu" - "${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_sparse_hdim128_fp16_sm80.cu" - "${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/flash_sparse_api.cpp" ) set(INCLUDES diff --git a/python/sglang/kernels/aot/benchmark/bench_awq_dequant.py b/python/sglang/kernels/aot/benchmark/bench_awq_dequant.py deleted file mode 100644 index cb22ba07f..000000000 --- a/python/sglang/kernels/aot/benchmark/bench_awq_dequant.py +++ /dev/null @@ -1,151 +0,0 @@ -import itertools -import os -from typing import List, Tuple - -import torch -import triton -import triton.testing -from sgl_kernel import awq_dequantize - -from sglang.utils import is_in_ci - -# Optional vLLM import -try: - from vllm import _custom_ops as ops - - VLLM_AVAILABLE = True -except ImportError: - ops = None - VLLM_AVAILABLE = False - -IS_CI = is_in_ci() - - -def vllm_awq_dequantize( - qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor -) -> Tuple[torch.Tensor, torch.Tensor]: - if not VLLM_AVAILABLE: - # Fallback to SGLang implementation - return sglang_awq_dequantize(qweight, scales, qzeros) - return ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0) - - -def sglang_awq_dequantize( - qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor -) -> Tuple[torch.Tensor, torch.Tensor]: - - return awq_dequantize(qweight, scales, qzeros) - - -def calculate_diff(qweight_row: int, qweight_col: int): - """Calculate difference between VLLM and SGLang implementations.""" - device = torch.device("cuda") - qweight = torch.randint( - 0, - torch.iinfo(torch.int32).max, - (qweight_row, qweight_col), - dtype=torch.int32, - device=device, - ) - group_size = qweight_row - scales_row = qweight_row // group_size - scales_col = qweight_col * 8 - scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device) - qzeros = torch.randint( - 0, - torch.iinfo(torch.int32).max, - (scales_row, qweight_col), - dtype=torch.int32, - device=device, - ) - - if not VLLM_AVAILABLE: - print("⚠️ vLLM not available, skipping comparison") - return - - vllm_out = vllm_awq_dequantize(qweight, scales, qzeros) - sglang_out = sglang_awq_dequantize(qweight, scales, qzeros) - - output_diff = torch.abs(vllm_out.float() - sglang_out.float()).mean().item() - - if torch.allclose( - vllm_out.to(torch.float32), sglang_out.to(torch.float32), rtol=1e-3, atol=1e-5 - ): - print("✅ All implementations match") - else: - print("❌ Implementations differ") - - -# CI environment uses simplified parameters -if IS_CI: - qweight_row_range = [128] # Single row size for CI - qweight_cols_range = [16] # Single column size for CI -else: - qweight_row_range = [3584, 18944, 128, 256, 512, 1024] - qweight_cols_range = [448, 576, 4736, 16, 32, 64, 128] - -configs = list(itertools.product(qweight_row_range, qweight_cols_range)) - - -@triton.testing.perf_report( - triton.testing.Benchmark( - x_names=["qweight_row", "qweight_col"], - x_vals=configs, - line_arg="provider", - line_vals=["vllm", "sglang"] if VLLM_AVAILABLE else ["sglang"], - line_names=["VLLM", "SGL Kernel"] if VLLM_AVAILABLE else ["SGL Kernel"], - styles=[("blue", "-"), ("green", "-")] if VLLM_AVAILABLE else [("green", "-")], - ylabel="us", - plot_name="awq-dequantize-performance", - args={}, - ) -) -def benchmark(qweight_row, qweight_col, provider): - dtype = torch.float16 - device = torch.device("cuda") - qweight = torch.randint( - 0, - torch.iinfo(torch.int32).max, - (qweight_row, qweight_col), - dtype=torch.int32, - device=device, - ) - group_size = qweight_row - scales_row = qweight_row // group_size - scales_col = qweight_col * 8 - scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device) - qzeros = torch.randint( - 0, - torch.iinfo(torch.int32).max, - (scales_row, qweight_col), - dtype=torch.int32, - device=device, - ) - - quantiles = [0.5, 0.2, 0.8] - - if provider == "vllm": - if not VLLM_AVAILABLE: - return (0, 0, 0) - fn = lambda: vllm_awq_dequantize( - qweight.clone(), scales.clone(), qzeros.clone() - ) - elif provider == "sglang": - fn = lambda: sglang_awq_dequantize( - qweight.clone(), scales.clone(), qzeros.clone() - ) - - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles) - - return 1000 * ms, 1000 * max_ms, 1000 * min_ms - - -if __name__ == "__main__": - # Simplify for CI environment - if IS_CI: - qweight_row, qweight_col = 128, 16 # Smaller values for CI - else: - qweight_row, qweight_col = 3584, 448 - - calculate_diff(qweight_row=qweight_row, qweight_col=qweight_col) - benchmark.run(print_data=True) diff --git a/python/sglang/kernels/aot/benchmark/bench_cutlass_mla.py b/python/sglang/kernels/aot/benchmark/bench_cutlass_mla.py deleted file mode 100644 index 9d24443f6..000000000 --- a/python/sglang/kernels/aot/benchmark/bench_cutlass_mla.py +++ /dev/null @@ -1,173 +0,0 @@ -import argparse -import copy -import itertools -import os - -import torch -import triton -from sgl_kernel import cutlass_mla_decode, cutlass_mla_get_workspace_size - -from sglang.srt.utils import get_device_capability -from sglang.utils import is_in_ci - -IS_CI = is_in_ci() - -# CI environment uses simplified parameters -if IS_CI: - bs_range = [1] # Single batch size for CI - qlen_range = [64] # Single sequence length for CI -else: - bs_range = [1, 8, 32, 64, 128, 256] - qlen_range = [1, 64, 128, 256, 512, 1024, 2048, 4096, 8192] - -configs = list(itertools.product(bs_range, qlen_range)) - - -@triton.testing.perf_report( - triton.testing.Benchmark( - x_names=["batch_size", "seq_len"], - x_vals=configs, - x_log=False, - line_arg="provider", - line_vals=[ - "128 heads", - "64 heads", - "32 heads", - "16 heads", - ], - line_names=[ - "128 heads", - "64 heads", - "32 heads", - "16 heads", - ], - styles=[("green", "-"), ("green", "--"), ("blue", "-"), ("blue", "--")], - ylabel="GB/s", - plot_name="cutlass mla", - args={}, - ) -) -def benchmark(batch_size, seq_len, provider, block_size, num_kv_splits): - d = 576 - dn = 64 - dv = 512 - - h_q_map = { - "128": 128, - "64": 64, - "32": 32, - "16": 16, - } - parsed_h_q = next( - (value for key, value in h_q_map.items() if key in provider), None - ) - - if parsed_h_q is None: - raise ValueError(f"Unknown head configuration in provider: {provider}") - h_q = parsed_h_q - - seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda") - max_seq_len = seq_lens.max().item() - block_num = (max_seq_len + block_size - 1) // block_size - - # Pad block_num so that small blocks can be packed into full 128-sized CUTLASS tiles. - # One 128-wide tile can hold (128 // block_size) small blocks. - pack_factor = 128 // block_size - block_num = ((block_num + pack_factor - 1) // pack_factor) * pack_factor - - qn = ( - torch.randn(h_q, batch_size, d - dn, dtype=torch.bfloat16, device="cuda") - * 100.0 - ) - qr = torch.randn(batch_size, h_q, dn, dtype=torch.bfloat16, device="cuda") * 100.0 - block_table = torch.randint( - 0, - batch_size * block_num, - (batch_size, block_num), - dtype=torch.int32, - device="cuda", - ) - - kv_cache = torch.randn( - block_table.numel(), block_size, d, dtype=torch.bfloat16, device="cuda" - ) - - workspace_size = cutlass_mla_get_workspace_size( - block_num * block_size, batch_size, num_kv_splits=num_kv_splits - ) - workspace = torch.empty(workspace_size, device="cuda", dtype=torch.uint8) - - quantiles = [0.5, 0.2, 0.8] - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( - lambda: cutlass_mla_decode( - qn.transpose(0, 1), - qr, - kv_cache, - seq_lens, - block_table, - workspace, - 1.44, - num_kv_splits, - ), - quantiles=quantiles, - ) - - q_size = qn.numel() * qn.element_size() + qr.numel() * qr.element_size() - - gbps = lambda ms: ( - (q_size + q_size * dv / d + kv_cache.numel() * kv_cache.element_size()) - * 1e-9 - / (ms * 1e-3) - ) - return gbps(ms), gbps(max_ms), gbps(min_ms) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - "--block-sizes", - nargs="+", - type=int, - default=[1, 32, 64, 128], - help="List of batch sizes", - ) - parser.add_argument( - "--num-kv-splits", - nargs="+", - type=int, - default=[-1], - help="List of batch sizes", - ) - args = parser.parse_args() - - # Skip in CI environment or unsupported architectures - if IS_CI: - major, minor = get_device_capability() - if major is None or major < 10: # Requires compute capability 10.0+ - print("Skipping Cutlass MLA benchmark in CI environment") - if major is not None: - print( - f"Cutlass MLA requires compute capability 10.0+, but found {major}.{minor}" - ) - else: - print("Could not determine device capability") - else: - for block_size in args.block_sizes: - for kv_split in args.num_kv_splits: - print(f"block_size={block_size}, num_kv_splits={kv_split}: ") - benchmark.run( - print_data=True, - block_size=block_size, - num_kv_splits=kv_split, - ) - print("Benchmark finished!") - else: - for block_size in args.block_sizes: - for kv_split in args.num_kv_splits: - print(f"block_size={block_size}, num_kv_splits={kv_split}: ") - benchmark.run( - print_data=True, - block_size=block_size, - num_kv_splits=kv_split, - ) - print("Benchmark finished!") diff --git a/python/sglang/kernels/aot/csrc/attention/cutlass_mla_kernel.cu b/python/sglang/kernels/aot/csrc/attention/cutlass_mla_kernel.cu deleted file mode 100644 index a41779c1b..000000000 --- a/python/sglang/kernels/aot/csrc/attention/cutlass_mla_kernel.cu +++ /dev/null @@ -1,274 +0,0 @@ -/* -Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -Copyright 2025 SGLang Team. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include -#include -#include -#include -#include - -#include -#include - -#include "cutlass_sm100_mla/device/sm100_mla.hpp" -#include "cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp" -#include "utils.h" - -// clang-format off -#if !defined(CUDA_VERSION) || CUDA_VERSION < 12040 -void cutlass_mla_decode( - torch::Tensor const& out, - torch::Tensor const& q_nope, - torch::Tensor const& q_pe, - torch::Tensor const& kv_c_and_k_pe_cache, - torch::Tensor const& seq_lens, - torch::Tensor const& page_table, - torch::Tensor const& workspace, - int64_t num_kv_splits) { - TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_decode"); -} -int64_t cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_batches, int64_t sm_count, int64_t num_kv_splits) { - TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_get_workspace_size"); -} -#else - -#define CUTLASS_CHECK(status) \ - { \ - cutlass::Status error = status; \ - TORCH_CHECK(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \ - } - -using namespace cute; -using namespace cutlass::fmha::kernel; - -template -struct IsPersistent { - static const bool value = v; -}; - -template > -struct MlaSm100 { - using Element = T; - using ElementAcc = float; - using ElementOut = T; - - using TileShape = Shape<_128, _128, Shape<_512, _64>>; - using TileShapeH = cute::tuple_element_t<0, TileShape>; - using TileShapeD = cute::tuple_element_t<2, TileShape>; - - // H K (D_latent D_rope) B - using ProblemShape = cute::tuple; - - using StrideQ = cute::tuple; // H D B - using StrideK = cute::tuple; // K D B - using StrideO = StrideK; // H D B - using StrideLSE = cute::tuple<_1, int>; // H B - - using TileScheduler = - std::conditional_t; - - using FmhaKernel = cutlass::fmha::kernel::Sm100FmhaMlaKernelTmaWarpspecialized< - TileShape, - Element, - ElementAcc, - ElementOut, - ElementAcc, - TileScheduler, - /*kIsCpAsync=*/!IsPaged128>; - using Fmha = cutlass::fmha::device::MLA; -}; - -template -typename T::Fmha::Arguments args_from_options( - at::Tensor const& out, - at::Tensor const& q_nope, - at::Tensor const& q_pe, - at::Tensor const& kv_c_and_k_pe_cache, - at::Tensor const& seq_lens, - at::Tensor const& page_table, - double sm_scale, - int64_t num_kv_splits) { - cutlass::KernelHardwareInfo hw_info; - hw_info.device_id = q_nope.device().index(); - hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); - - int batches = q_nope.size(0); - int page_count_per_seq = page_table.size(1); - int page_count_total = kv_c_and_k_pe_cache.size(0); - int page_size = kv_c_and_k_pe_cache.size(1); - int max_seq_len = page_size * page_count_per_seq; - using TileShapeH = typename T::TileShapeH; - using TileShapeD = typename T::TileShapeD; - auto problem_shape = cute::make_tuple(TileShapeH{}, max_seq_len, TileShapeD{}, batches); - - auto [H, K, D, B] = problem_shape; - auto [D_latent, D_rope] = D; - - float scale = float(sm_scale); - - using StrideQ = typename T::StrideQ; - using StrideK = typename T::StrideK; - using StrideO = typename T::StrideO; - using StrideLSE = typename T::StrideLSE; - - StrideQ stride_Q_nope = cute::make_tuple( - static_cast(q_nope.stride(1)), _1{}, static_cast(q_nope.stride(0))); - StrideQ stride_Q_pe = cute::make_tuple( - static_cast(q_pe.stride(1)), _1{}, static_cast(q_pe.stride(0))); - - StrideK stride_C = cute::make_tuple( - static_cast(0 + D_latent + D_rope), _1{}, static_cast(page_size * (D_latent + D_rope))); - StrideLSE stride_PT = cute::make_stride(_1{}, page_count_per_seq); - StrideLSE stride_LSE = cute::make_tuple(_1{}, 0 + H); - StrideO stride_O = cute::make_tuple(static_cast(0 + D_latent), _1{}, static_cast(0 + H * D_latent)); - - using Element = typename T::Element; - using ElementOut = typename T::ElementOut; - using ElementAcc = typename T::ElementAcc; - auto Q_nope_ptr = static_cast(q_nope.data_ptr()); - auto Q_pe_ptr = static_cast(q_pe.data_ptr()); - auto C_ptr = static_cast(kv_c_and_k_pe_cache.data_ptr()); - typename T::Fmha::Arguments arguments{ - problem_shape, - {scale, - Q_nope_ptr, - stride_Q_nope, - Q_pe_ptr, - stride_Q_pe, - C_ptr, - stride_C, - C_ptr + D_latent, - stride_C, - static_cast(seq_lens.data_ptr()), - static_cast(page_table.data_ptr()), - stride_PT, - page_count_total, - page_size}, - {static_cast(out.data_ptr()), stride_O, static_cast(nullptr), stride_LSE}, - hw_info, - // TODO(trevor-m): Change split_kv back to -1 when - // https://github.com/NVIDIA/cutlass/issues/2274 is fixed. Split_kv=1 will - // perform worse with larger context length and smaller batch sizes. - static_cast(num_kv_splits), // split_kv - nullptr, // is_var_split_kv - }; - // TODO(kaixih@nvidia): When split_kv=-1 and is_var_split_kv=false, we compute - // split_kv automatically based on batch size and sequence length to balance - // workload across available SMs. Consider using var_split_kv for manual - // control if needed. - T::Fmha::set_split_kv(arguments); - return arguments; -} - -template -void runMla( - at::Tensor const& out, - at::Tensor const& q_nope, - at::Tensor const& q_pe, - at::Tensor const& kv_c_and_k_pe_cache, - at::Tensor const& seq_lens, - at::Tensor const& page_table, - at::Tensor const& workspace, - double sm_scale, - int64_t num_kv_splits, - cudaStream_t stream) { - using MlaSm100Type = MlaSm100; - typename MlaSm100Type::Fmha fmha; - auto arguments = args_from_options(out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, sm_scale, num_kv_splits); - - CUTLASS_CHECK(fmha.can_implement(arguments)); - - CUTLASS_CHECK(fmha.initialize(arguments, workspace.data_ptr(), stream)); - - CUTLASS_CHECK(fmha.run(arguments, workspace.data_ptr(), stream)); -} - -#define DISPATCH_BOOL(expr, const_expr, ...) \ - [&]() -> bool { \ - if (expr) { \ - constexpr bool const_expr = true; \ - return __VA_ARGS__(); \ - } else { \ - constexpr bool const_expr = false; \ - return __VA_ARGS__(); \ - } \ - }() - -void cutlass_mla_decode( - torch::Tensor const& out, - torch::Tensor const& q_nope, - torch::Tensor const& q_pe, - torch::Tensor const& kv_c_and_k_pe_cache, - torch::Tensor const& seq_lens, - torch::Tensor const& page_table, - torch::Tensor const& workspace, - double sm_scale, - int64_t num_kv_splits) { - auto sm_version = getSMVersion(); - // On SM103a, half of the accuracy tests are failing. - TORCH_CHECK(sm_version == 100, "cutlass_mla_decode is only supported on compute capability 10.0, but found sm version ", sm_version); - - auto in_dtype = q_nope.dtype(); - at::cuda::CUDAGuard device_guard{(char)q_nope.get_device()}; - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(q_nope.get_device()); - const int page_size = kv_c_and_k_pe_cache.size(1); - - // NOTE(alcanderian): IsPersistent has bug with manual split_kv. - // Kernel will hang if batch is too large with large num_kv_splits. (for example bs=8, num_kv_splits=8) - // Maybe per batch split kv will fix this. - DISPATCH_BOOL(page_size == 128, IsPaged128, [&] { - DISPATCH_BOOL(num_kv_splits <= 1, NotManualSplitKV, [&] { - if (in_dtype == at::ScalarType::Half) { - runMla>( - out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream); - } else if (in_dtype == at::ScalarType::BFloat16) { - runMla>( - out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream); - } else if (in_dtype == at::ScalarType::Float8_e4m3fn) { - runMla>( - out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream); - } else { - TORCH_CHECK(false, "Unsupported input data type of MLA"); - } - return true; - }); - return true; - }); -} - -int64_t cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_batches, int64_t sm_count, int64_t num_kv_splits) { - // Workspace size depends on ElementAcc and ElementLSE (same as ElementAcc) - // which are float, so Element type here doesn't matter. - using MlaSm100Type = MlaSm100; - - // Get split kv. Requires problem shape and sm_count only. - typename MlaSm100Type::Fmha::Arguments arguments; - using TileShapeH = typename MlaSm100Type::TileShapeH; - using TileShapeD = typename MlaSm100Type::TileShapeD; - arguments.problem_shape = - cute::make_tuple(TileShapeH{}, static_cast(max_seq_len), TileShapeD{}, static_cast(num_batches)); - // Assumes device 0 when getting sm_count. - arguments.hw_info.sm_count = - sm_count <= 0 ? cutlass::KernelHardwareInfo::query_device_multiprocessor_count(/*device_id=*/0) : sm_count; - arguments.split_kv = static_cast(num_kv_splits); - MlaSm100Type::Fmha::set_split_kv(arguments); - - return MlaSm100Type::Fmha::get_workspace_size(arguments); -} - -#endif -// clang-format on diff --git a/python/sglang/kernels/aot/csrc/attention/cutlass_sm100_mla/device/sm100_mla.hpp b/python/sglang/kernels/aot/csrc/attention/cutlass_sm100_mla/device/sm100_mla.hpp deleted file mode 100644 index dd4ed231b..000000000 --- a/python/sglang/kernels/aot/csrc/attention/cutlass_sm100_mla/device/sm100_mla.hpp +++ /dev/null @@ -1,358 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ -/*! - \file - \brief An universal device layer for cutlass 3.x-style kernels. -*/ - -// clang-format off -#pragma once - -// common -#include "cutlass/cutlass.h" -#include "cutlass/device_kernel.h" - -#if !defined(__CUDACC_RTC__) -#include "cutlass/cluster_launch.hpp" -#include "cutlass/trace.h" -#endif // !defined(__CUDACC_RTC__) - -#include "../kernel/sm100_fmha_mla_tma_warpspecialized.hpp" -#include "../kernel/sm100_fmha_mla_reduction.hpp" - -//////////////////////////////////////////////////////////////////////////////// - -namespace cutlass::fmha::device { - -using namespace cute; -using namespace cutlass::fmha::kernel; - - -//////////////////////////////////////////////////////////////////////////////// -////////////////////////////// CUTLASS 3.x API ///////////////////////////////// -//////////////////////////////////////////////////////////////////////////////// - -template< - class Kernel_ -> -class MLA { -public: - - using Kernel = Kernel_; - - using ReductionKernel = cutlass::fmha::kernel::Sm100FmhaMlaReductionKernel< - typename Kernel::ElementOut, - typename Kernel::ElementAcc, - typename Kernel::ElementAcc, - Kernel::TileShapeH::value, - Kernel::TileShapeL::value, - 256 /*Max split*/ - >; - - /// Argument structure: User API - using KernelArguments = typename Kernel::Arguments; - using ReductionArguments = typename ReductionKernel::Arguments; - - using Arguments = KernelArguments; - - /// Argument structure: Kernel API - using KernelParams = typename Kernel::Params; - using ReductionParams = typename ReductionKernel::Params; - struct Params { - KernelParams fmha_params; - ReductionParams reduction_params; - }; - -private: - - /// Kernel API parameters object - Params params_; - - bool is_initialized(bool set = false) { - static bool initialized = false; - if (set) initialized = true; - return initialized; - } - - static ReductionArguments to_reduction_args(Arguments const& args) { - auto [H, K, D, B] = args.problem_shape; - return ReductionArguments{ - nullptr, args.epilogue.ptr_o, nullptr, args.epilogue.ptr_lse, - args.mainloop.softmax_scale, B, args.split_kv, K, args.mainloop.ptr_seq, - args.ptr_split_kv, Kernel::TileShapeS::value - }; - } - -public: - - /// Access the Params structure - Params const& params() const { - return params_; - } - - static void set_split_kv (KernelArguments& args) { - if (args.split_kv >= 1) return; - auto [H, K, D, B] = args.problem_shape; - int sm_count = args.hw_info.sm_count; - int max_splits = ceil_div(K, 128); - int sms_per_batch = max(1, sm_count / B); - int split_heur = min(max_splits, sms_per_batch); - int waves = ceil_div(B * split_heur, sm_count); - int k_waves = ceil_div(max_splits, split_heur); - int split_wave_aware = ceil_div(max_splits, k_waves); - args.split_kv = split_wave_aware; - } - - /// Determines whether the GEMM can execute the given problem. - static Status - can_implement(Arguments const& args) { - if (! Kernel::can_implement(args)) { - return Status::kInvalid; - } - if (! ReductionKernel::can_implement(to_reduction_args(args))) { - return Status::kInvalid; - } - return Status::kSuccess; - } - - /// Gets the workspace size - static size_t - get_workspace_size(Arguments const& args) { - size_t workspace_bytes = 0; - workspace_bytes += Kernel::get_workspace_size(args); - workspace_bytes += ReductionKernel::get_workspace_size(to_reduction_args(args)); - return workspace_bytes; - } - - /// Computes the maximum number of active blocks per multiprocessor - static int maximum_active_blocks(int /* smem_capacity */ = -1) { - CUTLASS_TRACE_HOST("MLA::maximum_active_blocks()"); - int max_active_blocks = -1; - int smem_size = Kernel::SharedStorageSize; - - // first, account for dynamic smem capacity if needed - cudaError_t result; - if (smem_size >= (48 << 10)) { - CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size); - result = cudaFuncSetAttribute( - device_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, - smem_size); - if (cudaSuccess != result) { - result = cudaGetLastError(); // to clear the error bit - CUTLASS_TRACE_HOST( - " cudaFuncSetAttribute() returned error: " - << cudaGetErrorString(result)); - return -1; - } - } - - // query occupancy after setting smem size - result = cudaOccupancyMaxActiveBlocksPerMultiprocessor( - &max_active_blocks, - device_kernel, - Kernel::MaxThreadsPerBlock, - smem_size); - - if (cudaSuccess != result) { - result = cudaGetLastError(); // to clear the error bit - CUTLASS_TRACE_HOST( - " cudaOccupancyMaxActiveBlocksPerMultiprocessor() returned error: " - << cudaGetErrorString(result)); - return -1; - } - - CUTLASS_TRACE_HOST(" max_active_blocks: " << max_active_blocks); - return max_active_blocks; - } - - /// Initializes GEMM state from arguments. - Status - initialize(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) { - CUTLASS_TRACE_HOST("MLA::initialize() - workspace " - << workspace << ", stream: " << (stream ? "non-null" : "null")); - - // Initialize the workspace - Status status = Kernel::initialize_workspace(args, workspace, stream); - if (status != Status::kSuccess) { - return status; - } - status = ReductionKernel::initialize_workspace(to_reduction_args(args), workspace, stream); - if (status != Status::kSuccess) { - return status; - } - KernelParams kernel_params = Kernel::to_underlying_arguments(args, workspace); - - ReductionArguments reduction_args = to_reduction_args(args); - if (reduction_args.split_kv > 1) { - reduction_args.ptr_oaccum = kernel_params.epilogue.ptr_o_acc; - reduction_args.ptr_lseaccum = kernel_params.epilogue.ptr_lse_acc; - } - ReductionParams reduction_params = ReductionKernel::to_underlying_arguments(reduction_args, workspace); - // Initialize the Params structure - params_ = Params {kernel_params, reduction_params}; - - if (is_initialized()) return Status::kSuccess; - - // account for dynamic smem capacity if needed - // no dynamic smem is needed for reduction kernel - int smem_size = Kernel::SharedStorageSize; - if (smem_size >= (48 << 10)) { - CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size); - cudaError_t result = cudaFuncSetAttribute( - device_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, - smem_size); - if (cudaSuccess != result) { - result = cudaGetLastError(); // to clear the error bit - CUTLASS_TRACE_HOST(" cudaFuncSetAttribute() returned error: " << cudaGetErrorString(result)); - return Status::kErrorInternal; - } - } - - is_initialized(true); - - return Status::kSuccess; - } - - /// Update API is preserved in 3.0, but does not guarantee a lightweight update of params. - Status - update(Arguments const& args, void* workspace = nullptr) { - CUTLASS_TRACE_HOST("MLA()::update() - workspace: " << workspace); - - size_t workspace_bytes = get_workspace_size(args); - if (workspace_bytes > 0 && nullptr == workspace) { - return Status::kErrorWorkspaceNull; - } - - auto fmha_params = Kernel::to_underlying_arguments(args, workspace); - - ReductionArguments reduction_args = to_reduction_args(args); - if (reduction_args.split_kv > 1) { - reduction_args.ptr_oaccum = fmha_params.epilogue.ptr_o_acc; - reduction_args.ptr_lseaccum = fmha_params.epilogue.ptr_lse_acc; - } - ReductionParams reduction_params = ReductionKernel::to_underlying_arguments(reduction_args, workspace); - // Initialize the Params structure - params_ = Params {fmha_params, reduction_params}; - - return Status::kSuccess; - } - - /// Primary run() entry point API that is static allowing users to create and manage their own params. - /// Supplied params struct must be construct by calling Kernel::to_underling_arguments() - static Status - run(Params& params, cudaStream_t stream = nullptr) { - CUTLASS_TRACE_HOST("MLA::run()"); - dim3 const block = Kernel::get_block_shape(); - dim3 const grid = Kernel::get_grid_shape(params.fmha_params); - - // configure smem size and carveout - int smem_size = Kernel::SharedStorageSize; - - Status launch_result; - // Use extended launch API only for mainloops that use it - if constexpr(Kernel::ArchTag::kMinComputeCapability >= 90) { - dim3 cluster(cute::size<0>(typename Kernel::ClusterShape{}), - cute::size<1>(typename Kernel::ClusterShape{}), - cute::size<2>(typename Kernel::ClusterShape{})); - void const* kernel = (void const*) device_kernel; - void* kernel_params[] = {¶ms.fmha_params}; - launch_result = ClusterLauncher::launch(grid, cluster, block, smem_size, stream, kernel, kernel_params); - } - else { - launch_result = Status::kSuccess; - device_kernel<<>>(params.fmha_params); - } - - cudaError_t result = cudaGetLastError(); - if (cudaSuccess != result or Status::kSuccess != launch_result) { - //return Status::kSuccess; - CUTLASS_TRACE_HOST(" Kernel launch failed. Reason: " << result); - return Status::kErrorInternal; - } - if (params.reduction_params.split_kv > 1) { - // launch reduction kernel - dim3 const block = ReductionKernel::get_block_shape(); - dim3 const grid = ReductionKernel::get_grid_shape(params.reduction_params); - device_kernel<<>>(params.reduction_params); - cudaError_t result = cudaGetLastError(); - if (cudaSuccess == result) { - return Status::kSuccess; - } - else { - CUTLASS_TRACE_HOST(" Kernel launch failed. Reason: " << result); - return Status::kErrorInternal; - } - } - else { - return Status::kSuccess; - } - } - - // - // Non-static launch overloads that first create and set the internal params struct of this kernel handle. - // - - /// Launches the kernel after first constructing Params internal state from supplied arguments. - Status - run(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) { - Status status = initialize(args, workspace, stream); - if (Status::kSuccess == status) { - status = run(params_, stream); - } - return status; - } - - /// Launches the kernel after first constructing Params internal state from supplied arguments. - Status - operator()(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) { - return run(args, workspace, stream); - } - - /// Overload that allows a user to re-launch the same kernel without updating internal params struct. - Status - run(cudaStream_t stream = nullptr) { - return run(params_, stream); - } - - /// Overload that allows a user to re-launch the same kernel without updating internal params struct. - Status - operator()(cudaStream_t stream = nullptr) { - return run(params_, stream); - } -}; - -//////////////////////////////////////////////////////////////////////////////// - -} // namespace cutlass::fmha::device - -//////////////////////////////////////////////////////////////////////////////// diff --git a/python/sglang/kernels/aot/csrc/attention/cutlass_sm100_mla/kernel/sm100_fmha_mla_reduction.hpp b/python/sglang/kernels/aot/csrc/attention/cutlass_sm100_mla/kernel/sm100_fmha_mla_reduction.hpp deleted file mode 100644 index b75870d0c..000000000 --- a/python/sglang/kernels/aot/csrc/attention/cutlass_sm100_mla/kernel/sm100_fmha_mla_reduction.hpp +++ /dev/null @@ -1,198 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// clang-format off -#pragma once - -#include "cutlass/cutlass.h" -#include "cutlass/arch/arch.h" -#include "cute/tensor.hpp" - -namespace cutlass::fmha::kernel { - -using namespace cute; -template< - class ElementOut, - class ElementAcc, - class ElementScale, - size_t kNumHeads, - size_t kHeadDimLatent, - int kMaxSplits -> -struct Sm100FmhaMlaReductionKernel { - - static const int SharedStorageSize = 0; - static const int MaxThreadsPerBlock = 128; - static const int MinBlocksPerMultiprocessor = 1; - - using ArchTag = cutlass::arch::Sm100; - - static_assert(kHeadDimLatent % MaxThreadsPerBlock == 0); - struct Arguments { - ElementAcc* ptr_oaccum = nullptr; - ElementOut* ptr_o = nullptr; - ElementAcc* ptr_lseaccum = nullptr; - ElementAcc* ptr_lse = nullptr; - ElementScale scale = 1.f; - int num_batches = 0; - int split_kv = -1; - int dim_k = -1; - int* ptr_seq = nullptr; - int* ptr_split_kv = nullptr; - int tile_shape_s = 128; - }; - using Params = Arguments; - - static Params to_underlying_arguments(Arguments const& args, void* workspace) { - return {args.ptr_oaccum, args.ptr_o, args.ptr_lseaccum, args.ptr_lse, - args.scale, args.num_batches, args.split_kv, args.dim_k, args.ptr_seq, - args.ptr_split_kv, args.tile_shape_s}; - } - - static size_t get_workspace_size(Arguments const& /*args*/) { - return 0; - } - - static Status initialize_workspace( - Arguments const& /*args*/, void* /*ws*/, cudaStream_t /*stream*/) { - return Status::kSuccess; - } - - static dim3 get_grid_shape(Params const& params) { - return dim3(kNumHeads, 1, params.num_batches); - } - - static dim3 get_block_shape() { - return dim3(MaxThreadsPerBlock, 1, 1); - } - - static bool can_implement(Arguments const& args) { - if (args.num_batches <= 0) return false; - if (args.split_kv <= 0) return false; - return true; - } - - CUTLASS_DEVICE void operator() (Params const& params, char* smem_raw) { - if (params.split_kv <= 1) return; - auto blk_coord = make_coord(blockIdx.x, _0{}, blockIdx.z); - - __shared__ ElementAcc sLseScale[kMaxSplits]; - const size_t offset_lseaccum = get<0>(blk_coord) + kNumHeads * params.split_kv * get<2>(blk_coord); - const size_t offset_lse = get<0>(blk_coord) + kNumHeads * get<2>(blk_coord); - - Tensor gLSEaccum = make_tensor(make_gmem_ptr(params.ptr_lseaccum + offset_lseaccum), - make_shape(params.split_kv), Stride>{}); - - Tensor gLSE = make_tensor(make_gmem_ptr(params.ptr_lse + offset_lse), - Shape<_1>{}, Stride<_1>{}); - - auto dim_k = params.ptr_seq == nullptr ? params.dim_k : params.ptr_seq[get<2>(blk_coord)]; - auto local_split_kv = params.ptr_split_kv == nullptr ? params.split_kv : params.ptr_split_kv[get<2>(blk_coord)]; - auto k_tile_total = ceil_div(dim_k, params.tile_shape_s); - auto k_tile_per_cta = ceil_div(k_tile_total, local_split_kv); - local_split_kv = ceil_div(k_tile_total, k_tile_per_cta); - - int warp_idx = cutlass::canonical_warp_idx_sync(); - if (warp_idx == 0) { - constexpr int kNLsePerThread = cute::ceil_div(kMaxSplits, 32); - - ElementAcc local_lse[kNLsePerThread]; - - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < kNLsePerThread; ++i) { - const int split = i * 32 + threadIdx.x; - local_lse[i] = split < local_split_kv ? gLSEaccum(split) : -std::numeric_limits::infinity(); - } - - ElementAcc lse_max = -std::numeric_limits::infinity(); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < kNLsePerThread; ++i) { - lse_max = max(lse_max, local_lse[i]); - } - CUTLASS_PRAGMA_UNROLL - for (int offset = 16; offset >= 1; offset /= 2) { - lse_max = max(lse_max, __shfl_xor_sync(0xffffffff, lse_max, offset)); - } - lse_max = lse_max == -std::numeric_limits::infinity() ? 0.0f : lse_max; // In case all local LSEs are -inf - lse_max = __shfl_sync(0xffffffff, lse_max, 0); - - ElementAcc sum_lse = 0; - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < kNLsePerThread; ++i) { - sum_lse = sum_lse + expf(local_lse[i] - lse_max); - } - - CUTLASS_PRAGMA_UNROLL - for (int offset = 16; offset >= 1; offset /= 2) { - sum_lse = sum_lse + __shfl_xor_sync(0xffffffff, sum_lse, offset); - } - - sum_lse = __shfl_sync(0xffffffff, sum_lse, 0); - - ElementAcc global_lse = (sum_lse == 0.f || sum_lse != sum_lse) ? std::numeric_limits::infinity() : logf(sum_lse) + lse_max; - if (threadIdx.x == 0 and params.ptr_lse != nullptr) { - gLSE(0) = global_lse; - } - - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < kNLsePerThread; ++i) { - const int split = i * 32 + threadIdx.x; - if (split < local_split_kv) { - sLseScale[split] = expf(local_lse[i] - global_lse); - } - } - } - __syncthreads(); - - constexpr int Elements = kHeadDimLatent / MaxThreadsPerBlock; - const size_t offset_oaccum = kHeadDimLatent * params.split_kv * (get<0>(blk_coord) + kNumHeads * get<2>(blk_coord)); - Tensor gOaccum = make_tensor(make_gmem_ptr(params.ptr_oaccum + offset_oaccum), - Shape>{}, Stride<_1>{}); - ElementAcc local_val[Elements] = {0}; - for (int split = 0; split < local_split_kv; ++split) { - ElementAcc lse_scale = sLseScale[split]; - CUTLASS_PRAGMA_UNROLL - for(int i = 0; i < Elements; ++i) { - local_val[i] += lse_scale * gOaccum(threadIdx.x + MaxThreadsPerBlock * i); - } - gOaccum.data() = gOaccum.data() + kHeadDimLatent; - } - auto ptr_o_local = params.ptr_o + (get<0>(blk_coord) + get<2>(blk_coord) * kNumHeads) * kHeadDimLatent; - Tensor gO = make_tensor(make_gmem_ptr(ptr_o_local), Shape>{}, Stride<_1>{}); - - CUTLASS_PRAGMA_UNROLL - for(int i = 0; i < Elements; ++i) { - gO(threadIdx.x + MaxThreadsPerBlock * i) = static_cast(local_val[i]); - } - } -}; - -} // namespace cutlass::fmha::kernel diff --git a/python/sglang/kernels/aot/csrc/attention/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp b/python/sglang/kernels/aot/csrc/attention/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp deleted file mode 100644 index 9809db84e..000000000 --- a/python/sglang/kernels/aot/csrc/attention/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp +++ /dev/null @@ -1,2018 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// clang-format off -#pragma once - -#include "cutlass/cutlass.h" - -#include "cute/tensor.hpp" -#include "cute/arch/simd_sm100.hpp" - -#include "cutlass/arch/arch.h" -#include "cutlass/arch/memory_sm80.h" -#include "cutlass/epilogue/thread/linear_combination.h" -#include "cutlass/gemm/collective/collective_builder.hpp" - -#include "gather_tensor.hpp" // from examples/common -#include "common/pow_2.hpp" - -namespace cutlass::fmha::kernel { - -using namespace cute; - -template< - class TileShape, - class Element_, - class ElementAcc_, - class ElementOut_, - class ElementLSE_, - class TileScheduler, -#ifdef CPASYNC - bool kIsCpAsync = true -#else - bool kIsCpAsync = false -#endif -> -struct Sm100FmhaMlaKernelTmaWarpspecialized { - - using Element = Element_; - using ElementAcc = ElementAcc_; - using ElementOut = ElementOut_; - using ElementLSE = ElementLSE_; - - // only 2Sm mode is supported - static const bool kIs2Sm = true; - static const int MaxThreadsPerBlock = 256; - static const int MinBlocksPerMultiprocessor = 1; - static const int TotalSNum = 2; - static const int TotalPNum = 2; - using ArchTag = cutlass::arch::Sm100; - - using ClusterShape = cute::conditional_t, Shape<_1, _1, _1>>; - - using TileShapeH = tuple_element_t<0, TileShape>; - using TileShapeS = tuple_element_t<1, TileShape>; - using TileShapeD = tuple_element_t<2, TileShape>; - - using TileShapeL = tuple_element_t<0, TileShapeD>; - using TileShapeR = tuple_element_t<1, TileShapeD>; - static_assert(TileShapeL{} % TileShapeR{} == 0, "Rope head dim must divide latent head dim"); - - using ProblemShape = Shape; - using TensorStride = Stride; - using TmemAllocator = cute::conditional_t; - - static_assert(TileShapeH{} == 128); - static const int kWarpsInN = kIs2Sm ? 2 : 1; - - static const int kNumComputeWarps = 4; - static const int kNumLoadWarps = kIsCpAsync ? 2 : 1; - - enum class WarpRole { - kMma = 0x1, kLoad = 0x2, kCompute = 0x3, kLoadPageTable = 0x4, kEmpty=0x0 - }; - - static const long long unsigned int kWarpAssignment = kIsCpAsync ? 0x4221'3333ull : 0x0021'3333ull; - - static CUTLASS_DEVICE WarpRole warp_idx_to_role(int warp_idx) { - return static_cast((kWarpAssignment >> (4 * warp_idx)) & 0xF); - } - - static const int Alignment = 128 / sizeof_bits_v; - static const int AlignmentOut = 128 / sizeof_bits_v; - - using TileShapeQK = Shape; - static const int StagesQK = 24 / sizeof(Element); // free parameter - static const int IterationsQKLatent = decltype(TileShapeL{} / get<2>(TileShapeQK{}))::value; - static const int IterationsQKRope = decltype(TileShapeR{} / get<2>(TileShapeQK{}))::value; - static const int IterationsQK = IterationsQKLatent + IterationsQKRope; - - using Schedule = cute::conditional_t; - using CollectiveMmaQK = typename cutlass::gemm::collective::CollectiveBuilder< - cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, - Element, TensorStride, Alignment, - Element, TensorStride, Alignment, - ElementAcc, - TileShapeQK, ClusterShape, cutlass::gemm::collective::StageCount, - Schedule>::CollectiveOp; - using TiledMmaQK = typename CollectiveMmaQK::TiledMma; - using CtaShapeQK = typename CollectiveMmaQK::CtaShape_MNK; - - // chosen for unified smem staging between K and V - using TileShapePV = Shape; - using TransposeTensorStride = decltype(select<1,0,2>(TensorStride{})); - static const int StagesPV = StagesQK; // not sure why, but must be at least two. check pipes - static const int IterationsPV_K = decltype(TileShapeS{} / get<2>(TileShapePV{}))::value; - static const int IterationsPV_N = decltype(TileShapeL{} / get<1>(TileShapePV{}))::value; - - using CollectiveMmaPV = typename cutlass::gemm::collective::CollectiveBuilder< - cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, - Element, TensorStride, Alignment, - Element, TransposeTensorStride, Alignment, - ElementAcc, - TileShapePV, ClusterShape, cutlass::gemm::collective::StageCount, - Schedule>::CollectiveOp; - using CtaShapePV = typename CollectiveMmaPV::CtaShape_MNK; - static_assert(std::is_same_v); - - using TiledMmaPV = typename CollectiveMmaPV::TiledMma; - - using AtomThrShapeMNK = typename CollectiveMmaQK::AtomThrShapeMNK; - static_assert(typename CollectiveMmaQK::AtomThrShapeMNK{} == typename CollectiveMmaPV::AtomThrShapeMNK{}, "schedule must match"); - - static const int StagesPageTable = kIsCpAsync ? StagesPV : 1; - - // pipelines from load to mma, PipelineTmaUmmaAsync, stages tbd - // use expect_tx for Q load - using PipelineLoadQK = cute::conditional_t, PipelineTmaUmmaAsync>; - using PipelineLoadPV = PipelineLoadQK; - // pipeline from mma (Q@K) to softmax, PipelineUmmaAsync, 2 stages - using PipelineS = PipelineUmmaAsync; - // pipeline from softmax (P) to mma (bmm2), PipelineUmmaAsync, 2 stages - using PipelineP = PipelineUmmaConsumerAsync; - // pipeline from mma to softmax (for rescale), PipelineUmmaAsync, 1 stage - using PipelineO = PipelineUmmaAsync<1, AtomThrShapeMNK>; - - using PipelinePT = PipelineAsync; - - struct PipelineStorage { - alignas(16) typename PipelineLoadQK::SharedStorage load_qk; - alignas(16) typename PipelineS::SharedStorage mma_s; - alignas(16) typename PipelineP::SharedStorage p_mma; - alignas(16) typename PipelineO::SharedStorage mma_o; - alignas(16) typename PipelinePT::SharedStorage load_page_table; - }; - - template - static CUTE_DEVICE constexpr auto unstageSmemLayout(Layout const& layout, Stages stages = {}) { - return composition(layout, make_tuple(_, _, _, make_layout(stages))); - } - - using SmemLayoutQ = decltype(unstageSmemLayout(typename CollectiveMmaQK::SmemLayoutA{}, Int{})); - using SmemLayoutKC = typename CollectiveMmaQK::SmemLayoutB; - using SmemLayoutVC = typename CollectiveMmaPV::SmemLayoutB; - using SmemLayoutP = decltype(unstageSmemLayout(typename CollectiveMmaPV::SmemLayoutA{}, make_shape(Int{}, _2{}))); - - static const int kBytesLoadQ = size(AtomThrShapeMNK{}) * cutlass::bits_to_bytes(cosize(take<0,3>(SmemLayoutQ{})) * cute::sizeof_bits_v); - static const int kBytesLoadKC = size(AtomThrShapeMNK{}) * cutlass::bits_to_bytes(cosize(take<0,3>(SmemLayoutKC{})) * cute::sizeof_bits_v); - static const int kBytesLoadVC = size(AtomThrShapeMNK{}) * cutlass::bits_to_bytes(cosize(take<0,3>(SmemLayoutVC{})) * cute::sizeof_bits_v); - // pre-condition for overlapped smem staging - static_assert(kBytesLoadKC == kBytesLoadVC); - static_assert(StagesQK == StagesPV); - - static const int kTransactionsBytesLoadQK = kBytesLoadKC; - static const int kTransactionsBytesLoadExtraQ = kBytesLoadQ; - static const int kTransactionsBytesLoadPV = kBytesLoadVC; - - static const int kNamedBarrierExchange = (int) cutlass::arch::ReservedNamedBarriers::TransformBarrier; - // This Named Barrier is introduced to solve Q tile loading overwritten issue when enable persistent - // tile scheduler for FP8 MLA. - static const int kNamedBarrierEpilogue = (int) cutlass::arch::ReservedNamedBarriers::EpilogueBarrier; - // - static const int kNamedBarrierTmemDealloc = (int) cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier; - - enum class TmemAllocation : uint32_t { - kSizeS = TileShapeS::value / kWarpsInN, - // Overall - kSizeO = TileShapeL::value / kWarpsInN, - // Between accumulators we loop over - kSizeAccO = decltype(get<1>(TileShapePV{}))::value / kWarpsInN, - kNumS = TotalSNum, - kNumP = TotalPNum, - kNumO = 1, - kS0 = 0, - kS1 = kS0 + kSizeS, - kO0 = kS1 + kSizeS, - kTotal = kO0 + kSizeO - }; - - static_assert(static_cast(TmemAllocation::kTotal) <= TmemAllocator::Sm100TmemCapacityColumns, "using too much tmem"); - - struct TensorStorage { - // to communicate max and row_sum - cute::array smem_exchange; - cute::array smem_page_table; - alignas(2048) cute::array> smem_q; - union { - alignas(2048) cute::array> smem_kc; - alignas(2048) cute::array> smem_vc; - }; - alignas(2048) cute::array> smem_p; - }; - - struct SharedStorage { - PipelineStorage pipelines; - TensorStorage tensors; - uint32_t tmem_base_ptr; - }; - - static const int SharedStorageSize = sizeof(SharedStorage); - static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "using too much smem"); - - struct MainloopArguments { - ElementAcc softmax_scale; - - // all tensors strides are (num_heads or seqlen, head_dim, batch) - // head_dim stride is always 1 - Element* ptr_q_latent; - TensorStride stride_q_latent; - Element* ptr_q_rope; - TensorStride stride_q_rope; - - Element* ptr_c_latent; - TensorStride stride_c_latent; - Element* ptr_k_rope; - TensorStride stride_k_rope; - - // for paged attention, we interpret what was previously [batch, seqlen] - // as [page_count, page_size], and index according to page_table - int* ptr_seq = nullptr; - int* ptr_page_table = nullptr; - // page table is [batch, seqlen or similar] - Stride<_1, int> stride_page_table = {}; - int page_count = 0; - int page_size = TileShapeS{}; // powers of two if kIsCpAsync, otherwise TileShapeS - }; - - struct EpilogueArguments { - ElementOut* ptr_o = nullptr; - TensorStride stride_o; - ElementLSE* ptr_lse = nullptr; - Stride<_1, int> stride_lse; - ElementAcc output_scale = 1.0f; - }; - - struct Arguments { - // (num_heads=128, seqlen, (d_latent=512, d_rope=64), batch_count) - // for paged attention, seqlen is max seqlen - ProblemShape problem_shape; - MainloopArguments mainloop; - EpilogueArguments epilogue; - KernelHardwareInfo hw_info; - int split_kv = -1; - int* ptr_split_kv = nullptr; - }; - - using TmaLoadQLatent = typename CollectiveMmaQK::Params::TMA_A; - using TmaLoadQRope = typename CollectiveMmaQK::Params::TMA_A; - using TmaLoadCLatent = typename CollectiveMmaQK::Params::TMA_B; - using TmaLoadKRope = typename CollectiveMmaQK::Params::TMA_B; - using TmaLoadCLatentTranspose = typename CollectiveMmaPV::Params::TMA_B; - - struct MainloopParams { - TmaLoadQLatent tma_load_q_latent; - TmaLoadQRope tma_load_q_rope; - TmaLoadCLatent tma_load_c_latent; - TmaLoadKRope tma_load_k_rope; - TmaLoadCLatentTranspose tma_load_c_latent_transpose; - }; - - struct EpilogueParams { - ElementOut* ptr_o = nullptr; - ElementAcc* ptr_o_acc = nullptr; - TensorStride stride_o; - TensorStride stride_o_acc; - ElementLSE* ptr_lse = nullptr; - ElementLSE* ptr_lse_acc = nullptr; - Stride<_1, int> stride_lse; - Stride<_1, int> stride_lse_acc; - ElementAcc output_scale = 1.0f; - }; - - struct Params { - ProblemShape problem_shape; - MainloopArguments mainloop; - EpilogueParams epilogue; - MainloopParams mainloop_params; - typename TileScheduler::Params tile_scheduler; - int split_kv = -1; - int* ptr_split_kv = nullptr; - }; - - static Params to_underlying_arguments(Arguments const& args, void* workspace) { - //workspace = nullptr; // let's get an error if one of these needs workspace - - auto [H, K, D, B] = args.problem_shape; - auto [L, R] = D; - - int paged_B = B; - int paged_K = K; - if (args.mainloop.ptr_page_table != nullptr) { - paged_B = args.mainloop.page_count; - paged_K = args.mainloop.page_size; - } - - auto params_qk_latent = CollectiveMmaQK::to_underlying_arguments( - make_shape(H, K, L, B), - typename CollectiveMmaQK::Arguments { - args.mainloop.ptr_q_latent, args.mainloop.stride_q_latent, - args.mainloop.ptr_c_latent, args.mainloop.stride_c_latent, - }, nullptr); - - auto params_qk_latent_paged = CollectiveMmaQK::to_underlying_arguments( - make_shape(H, paged_K, L, paged_B), - typename CollectiveMmaQK::Arguments { - args.mainloop.ptr_q_latent, args.mainloop.stride_q_latent, - args.mainloop.ptr_c_latent, args.mainloop.stride_c_latent, - }, nullptr); - - auto params_qk_rope = CollectiveMmaQK::to_underlying_arguments( - make_shape(H, K, R, B), - typename CollectiveMmaQK::Arguments { - args.mainloop.ptr_q_rope, args.mainloop.stride_q_rope, - args.mainloop.ptr_k_rope, args.mainloop.stride_k_rope, - }, nullptr); - - auto params_qk_rope_paged = CollectiveMmaQK::to_underlying_arguments( - make_shape(H, paged_K, R, paged_B), - typename CollectiveMmaQK::Arguments { - args.mainloop.ptr_q_rope, args.mainloop.stride_q_rope, - args.mainloop.ptr_k_rope, args.mainloop.stride_k_rope, - }, nullptr); - - - auto stride_c_latent_transpose = select<1,0,2>(args.mainloop.stride_c_latent); - auto params_pv_latent = CollectiveMmaPV::to_underlying_arguments( - make_shape(H, L, paged_K, paged_B), - typename CollectiveMmaPV::Arguments { - args.mainloop.ptr_q_latent, args.mainloop.stride_q_latent, // dummy, never used - args.mainloop.ptr_c_latent, stride_c_latent_transpose, - }, nullptr); - - MainloopParams mainloop_params { - params_qk_latent.tma_load_a, - params_qk_rope.tma_load_a, - params_qk_latent_paged.tma_load_b, - params_qk_rope_paged.tma_load_b, - params_pv_latent.tma_load_b - }; - - EpilogueParams epilogue_params; - - epilogue_params.ptr_o = args.epilogue.ptr_o; - epilogue_params.stride_o = args.epilogue.stride_o; - epilogue_params.ptr_lse = args.epilogue.ptr_lse; - epilogue_params.stride_lse = args.epilogue.stride_lse; - epilogue_params.output_scale = args.epilogue.output_scale; - - if (args.split_kv > 1) { - ElementAcc* ptr_o_acc = reinterpret_cast(workspace); - ElementLSE* ptr_lse_acc = reinterpret_cast(ptr_o_acc + H * L * args.split_kv * B); - epilogue_params.ptr_o_acc = ptr_o_acc; - epilogue_params.ptr_lse_acc = ptr_lse_acc; - - epilogue_params.stride_o_acc = make_tuple(static_cast(0 + L) * args.split_kv, _1{}, static_cast(0 + H * L) * args.split_kv); - epilogue_params.stride_lse_acc = make_tuple(_1{}, (0 + H) * args.split_kv); - } - - return {args.problem_shape, args.mainloop, epilogue_params, mainloop_params, - TileScheduler::to_underlying_arguments(args.problem_shape, args.hw_info, ClusterShape{}, args.split_kv), args.split_kv, args.ptr_split_kv}; - } - - static size_t get_workspace_size(Arguments const& args) { - ProblemShape problem_shape = args.problem_shape; - auto [H, K, D, B] = problem_shape; - auto [D_latent, D_rope] = D; - auto split_kv = args.split_kv; - return (sizeof(ElementAcc) * D_latent + sizeof(ElementLSE)) * H * split_kv * B; - } - static Status initialize_workspace( - Arguments const& /*args*/, void* /*ws*/, cudaStream_t /*stream*/) { - return Status::kSuccess; - } - - static dim3 get_grid_shape(Params const& params) { - return TileScheduler::get_grid_shape(params.tile_scheduler); - } - - static dim3 get_block_shape() { - dim3 block(MaxThreadsPerBlock, 1, 1); - return block; - } - - static bool can_implement(Arguments const& args) { - if (kIsCpAsync) { - if ((args.mainloop.page_size & (args.mainloop.page_size - 1)) != 0) { - return false; - } - if (args.mainloop.page_size > TileShapeS{}) { - return false; - } - } - else { - if (args.mainloop.ptr_page_table != nullptr && args.mainloop.page_size != TileShapeS{}) { - return false; - } - } - if (get<0>(args.problem_shape) != 128) { - return false; - } - if (get<1>(args.problem_shape) <= 0) { - return false; - } - if (args.split_kv <= 0) { - return false; - } - return true; - } - - - CUTLASS_DEVICE void operator()(Params const& params, char* smem_raw) { - - TileScheduler tile_scheduler(params.tile_scheduler); - - int warp_idx = cutlass::canonical_warp_idx_sync(); - auto role = warp_idx_to_role(warp_idx); - uint32_t lane_predicate = cute::elect_one_sync(); - - uint32_t cta_rank_in_cluster = cute::block_rank_in_cluster(); - int cta_coord_v = cta_rank_in_cluster % size<0>(AtomThrShapeMNK{}); - bool is_mma_leader_cta = cta_coord_v == 0; - - if (role == WarpRole::kLoad && lane_predicate && ! kIsCpAsync) { - prefetch_tma_descriptor(params.mainloop_params.tma_load_q_latent.get_tma_descriptor()); - prefetch_tma_descriptor(params.mainloop_params.tma_load_c_latent.get_tma_descriptor()); - prefetch_tma_descriptor(params.mainloop_params.tma_load_q_rope.get_tma_descriptor()); - prefetch_tma_descriptor(params.mainloop_params.tma_load_k_rope.get_tma_descriptor()); - prefetch_tma_descriptor(params.mainloop_params.tma_load_c_latent_transpose.get_tma_descriptor()); - } - SharedStorage& shared_storage = *reinterpret_cast(smem_raw); - - typename PipelineLoadQK::Params pipeline_load_qk_params; - if (role == WarpRole::kLoad) { - pipeline_load_qk_params.role = PipelineLoadQK::ThreadCategory::Producer; - } - if (role == WarpRole::kMma) { - pipeline_load_qk_params.role = PipelineLoadQK::ThreadCategory::Consumer; - } - if constexpr (kIsCpAsync) { - // we can make our life easier by unconditionally loading blocks - // since we know it'll always be legal - pipeline_load_qk_params.producer_arv_count = kNumLoadWarps * cutlass::NumThreadsPerWarp * size(AtomThrShapeMNK{}); - } - else { - pipeline_load_qk_params.is_leader = lane_predicate && (role == WarpRole::kLoad) && is_mma_leader_cta; - pipeline_load_qk_params.transaction_bytes = kTransactionsBytesLoadQK; - } - pipeline_load_qk_params.initializing_warp = 0; - PipelineLoadQK pipeline_load_qk(shared_storage.pipelines.load_qk, pipeline_load_qk_params, - ClusterShape{}, /*barrier init*/ cute::true_type{}, /*mask calc*/cute::false_type{}); - - typename PipelineS::Params pipeline_mma_s_params; - if (role == WarpRole::kMma) { - pipeline_mma_s_params.role = PipelineS::ThreadCategory::Producer; - } - if (role == WarpRole::kCompute) { - pipeline_mma_s_params.role = PipelineS::ThreadCategory::Consumer; - } - pipeline_mma_s_params.consumer_arv_count = kNumComputeWarps * cutlass::NumThreadsPerWarp * size(AtomThrShapeMNK{}); - pipeline_mma_s_params.initializing_warp = 1; - PipelineS pipeline_mma_s( - shared_storage.pipelines.mma_s, - pipeline_mma_s_params, - ClusterShape{}, /*barrier init*/ cute::true_type{}, /*mask calc*/cute::false_type{}); - - typename PipelineP::Params pipeline_p_mma_params; - if (role == WarpRole::kMma) { - pipeline_p_mma_params.role = PipelineP::ThreadCategory::Consumer; - } - if (role == WarpRole::kCompute) { - pipeline_p_mma_params.role = PipelineP::ThreadCategory::Producer; - } - pipeline_p_mma_params.producer_arv_count = kNumComputeWarps * cutlass::NumThreadsPerWarp * size(AtomThrShapeMNK{}); - pipeline_p_mma_params.consumer_arv_count = 1; - pipeline_p_mma_params.initializing_warp = 2; - PipelineP pipeline_p_mma( - shared_storage.pipelines.p_mma, - pipeline_p_mma_params, - ClusterShape{}, /*barrier init*/ cute::true_type{}, /*mask calc*/cute::false_type{}); - - typename PipelineO::Params pipeline_mma_o_params; - if (role == WarpRole::kMma) { - pipeline_mma_o_params.role = PipelineO::ThreadCategory::Producer; - } - if (role == WarpRole::kCompute) { - pipeline_mma_o_params.role = PipelineO::ThreadCategory::Consumer; - } - pipeline_mma_o_params.consumer_arv_count = kNumComputeWarps * cutlass::NumThreadsPerWarp * size(AtomThrShapeMNK{}); - pipeline_mma_o_params.initializing_warp = 3; - PipelineO pipeline_mma_o( - shared_storage.pipelines.mma_o, - pipeline_mma_o_params, - ClusterShape{}, /*barrier init*/ cute::true_type{}, /*mask calc*/cute::false_type{}); - - typename PipelinePT::Params pipeline_pt_params; - if (role == WarpRole::kLoad) { - pipeline_pt_params.role = PipelinePT::ThreadCategory::Consumer; - } - if (role == WarpRole::kLoadPageTable) { - pipeline_pt_params.role = PipelinePT::ThreadCategory::Producer; - } - pipeline_pt_params.consumer_arv_count = kNumLoadWarps * cutlass::NumThreadsPerWarp; - pipeline_pt_params.producer_arv_count = cutlass::NumThreadsPerWarp; - pipeline_pt_params.initializing_warp = 4; - PipelinePT pipeline_page_table( - shared_storage.pipelines.load_page_table, - pipeline_pt_params); - - TmemAllocator tmem_allocator; - - pipeline_init_arrive_relaxed(size(ClusterShape{})); - - pipeline_load_qk.init_masks(ClusterShape{}); // do we need an update here for 2Sm? - pipeline_mma_s.init_masks(ClusterShape{}); - pipeline_p_mma.init_masks(ClusterShape{}); - pipeline_mma_o.init_masks(ClusterShape{}); - - typename PipelineLoadQK::PipelineState pipeline_load_qk_consumer_state; - typename PipelineLoadQK::PipelineState pipeline_load_qk_producer_state = cutlass::make_producer_start_state(); - - typename PipelineS::PipelineState pipeline_mma_s_consumer_state; - typename PipelineS::PipelineState pipeline_mma_s_producer_state = cutlass::make_producer_start_state(); - - typename PipelineP::PipelineState pipeline_p_mma_consumer_state; - typename PipelineP::PipelineState pipeline_p_mma_producer_state = cutlass::make_producer_start_state(); - - typename PipelineO::PipelineState pipeline_mma_o_consumer_state; - typename PipelineO::PipelineState pipeline_mma_o_producer_state = cutlass::make_producer_start_state(); - - typename PipelinePT::PipelineState pipeline_pt_consumer_state; - typename PipelinePT::PipelineState pipeline_pt_producer_state = cutlass::make_producer_start_state(); - - pipeline_init_wait(size(ClusterShape{})); - - if (role == WarpRole::kLoadPageTable) { - CUTLASS_PRAGMA_NO_UNROLL - for (; tile_scheduler.is_valid(); ++tile_scheduler) { - auto blk_coord = tile_scheduler.get_block_coord(); - auto problem_shape = params.problem_shape; - auto local_split_kv = params.split_kv; - if (params.mainloop.ptr_seq != nullptr) { - get<1>(problem_shape) = params.mainloop.ptr_seq[get<2>(blk_coord)]; - if (params.ptr_split_kv != nullptr) { - local_split_kv = params.ptr_split_kv[get<2>(blk_coord)]; - } - } - if (local_split_kv <= get<3>(blk_coord)) - continue; - load_page_table( - blk_coord, - problem_shape, - params.mainloop, - shared_storage.tensors, - pipeline_page_table, pipeline_pt_producer_state, - local_split_kv - ); - } - } - else if (role == WarpRole::kLoad) { - if constexpr (kIsCpAsync) { - CUTLASS_PRAGMA_NO_UNROLL - for (; tile_scheduler.is_valid(); ++tile_scheduler) { - auto blk_coord = tile_scheduler.get_block_coord(); - auto problem_shape = params.problem_shape; - auto local_split_kv = params.split_kv; - if (params.mainloop.ptr_seq != nullptr) { - get<1>(problem_shape) = params.mainloop.ptr_seq[get<2>(blk_coord)]; - if (params.ptr_split_kv != nullptr) { - local_split_kv = params.ptr_split_kv[get<2>(blk_coord)]; - } - } - if (local_split_kv <= get<3>(blk_coord)) - continue; - load_cpasync( - blk_coord, - problem_shape, - params.mainloop, - params.mainloop_params, - shared_storage.tensors, - pipeline_load_qk, pipeline_load_qk_producer_state, - local_split_kv, - /* must be shared pipe */ - pipeline_page_table, pipeline_pt_consumer_state - ); - cutlass::arch::NamedBarrier((kNumComputeWarps + kNumLoadWarps) * NumThreadsPerWarp, kNamedBarrierEpilogue).arrive_and_wait(); - } - } - else { - if (params.mainloop.ptr_page_table != nullptr) { - CUTLASS_PRAGMA_NO_UNROLL - for (; tile_scheduler.is_valid(); ++tile_scheduler) { - auto blk_coord = tile_scheduler.get_block_coord(); - auto problem_shape = params.problem_shape; - auto local_split_kv = params.split_kv; - if (params.mainloop.ptr_seq != nullptr) { - get<1>(problem_shape) = params.mainloop.ptr_seq[get<2>(blk_coord)]; - if (params.ptr_split_kv != nullptr) { - local_split_kv = params.ptr_split_kv[get<2>(blk_coord)]; - } - } - if (local_split_kv <= get<3>(blk_coord)) - continue; - load_tma( - blk_coord, - problem_shape, - params.mainloop, - params.mainloop_params, - shared_storage.tensors, - pipeline_load_qk, pipeline_load_qk_producer_state, - pipeline_load_qk, pipeline_load_qk_producer_state, - local_split_kv - ); - cutlass::arch::NamedBarrier((kNumComputeWarps + kNumLoadWarps) * NumThreadsPerWarp, kNamedBarrierEpilogue).arrive_and_wait(); - } - } - else { - CUTLASS_PRAGMA_NO_UNROLL - for (; tile_scheduler.is_valid(); ++tile_scheduler) { - auto blk_coord = tile_scheduler.get_block_coord(); - auto problem_shape = params.problem_shape; - auto local_split_kv = params.split_kv; - if (params.mainloop.ptr_seq != nullptr) { - get<1>(problem_shape) = params.mainloop.ptr_seq[get<2>(blk_coord)]; - if (params.ptr_split_kv != nullptr) { - local_split_kv = params.ptr_split_kv[get<2>(blk_coord)]; - } - } - if (local_split_kv <= get<3>(blk_coord)) - continue; - load_tma( - blk_coord, - problem_shape, - params.mainloop, - params.mainloop_params, - shared_storage.tensors, - pipeline_load_qk, pipeline_load_qk_producer_state, - pipeline_load_qk, pipeline_load_qk_producer_state, - local_split_kv - ); - cutlass::arch::NamedBarrier((kNumComputeWarps + kNumLoadWarps) * NumThreadsPerWarp, kNamedBarrierEpilogue).arrive_and_wait(); - } - } - } - } - else if (role == WarpRole::kMma) { - tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr); - __syncwarp(); - - if (is_mma_leader_cta) { - CUTLASS_PRAGMA_NO_UNROLL - for (; tile_scheduler.is_valid(); ++tile_scheduler) { - auto blk_coord = tile_scheduler.get_block_coord(); - auto problem_shape = params.problem_shape; - auto local_split_kv = params.split_kv; - if (params.mainloop.ptr_seq != nullptr) { - get<1>(problem_shape) = params.mainloop.ptr_seq[get<2>(blk_coord)]; - if (params.ptr_split_kv != nullptr) { - local_split_kv = params.ptr_split_kv[get<2>(blk_coord)]; - } - } - if (local_split_kv <= get<3>(blk_coord)) - continue; - mma(blk_coord, - problem_shape, - shared_storage.tensors, - pipeline_load_qk, pipeline_load_qk_consumer_state, - pipeline_load_qk, pipeline_load_qk_consumer_state, - pipeline_mma_s, pipeline_mma_s_producer_state, - pipeline_p_mma, pipeline_p_mma_consumer_state, - pipeline_mma_o, pipeline_mma_o_producer_state, - local_split_kv - ); - } - } - - //cutlass::arch::NamedBarrier((kNumComputeWarps + 1) * NumThreadsPerWarp, kNamedBarrierTmemDealloc).arrive_and_wait(); - - //uint32_t free_stage_ptr = shared_storage.tmem_base_ptr; - //tmem_allocator.free(free_stage_ptr, TmemAllocator::Sm100TmemCapacityColumns); - } - else if (role == WarpRole::kCompute) { - CUTLASS_PRAGMA_NO_UNROLL - for (; tile_scheduler.is_valid(); ++tile_scheduler) { - auto blk_coord = tile_scheduler.get_block_coord(); - auto problem_shape = params.problem_shape; - auto split_kv = params.split_kv; - auto local_split_kv = split_kv; - if (params.mainloop.ptr_seq != nullptr) { - get<1>(problem_shape) = params.mainloop.ptr_seq[get<2>(blk_coord)]; - if (params.ptr_split_kv != nullptr) { - local_split_kv = params.ptr_split_kv[get<2>(blk_coord)]; - } - } - if (local_split_kv <= get<3>(blk_coord)) - continue; - compute( - blk_coord, - problem_shape, - params.mainloop, // for softmax_scale - params.epilogue, - shared_storage.tensors, // for smem_comm - pipeline_mma_s, pipeline_mma_s_consumer_state, - pipeline_p_mma, pipeline_p_mma_producer_state, - pipeline_mma_o, pipeline_mma_o_consumer_state, - local_split_kv - ); - } - - //cutlass::arch::NamedBarrier((kNumComputeWarps + 1) * NumThreadsPerWarp, kNamedBarrierTmemDealloc).arrive(); - } - - cute::cluster_sync(); - cutlass::arch::NamedBarrier((kNumComputeWarps + 1) * NumThreadsPerWarp, kNamedBarrierTmemDealloc).arrive(); - if (role == WarpRole::kMma) { - uint32_t free_stage_ptr = shared_storage.tmem_base_ptr; - tmem_allocator.free(free_stage_ptr, TmemAllocator::Sm100TmemCapacityColumns); - } - } - - template - CUTLASS_DEVICE void load_page_table( - BlkCoord const& blk_coord, - ProblemShape const& problem_shape, - MainloopArguments const& mainloop_args, - TensorStorage& shared_tensors, - PipelinePT& pipeline_page_table, - typename PipelinePT::PipelineState& pipeline_pt_producer_state, int const& split_kv) { - - auto [H, K, D, B] = problem_shape; - int batch_coord = get<2>(blk_coord); - - auto mPT_l = make_tensor(make_gmem_ptr(mainloop_args.ptr_page_table), - make_shape(mainloop_args.page_count, B), - mainloop_args.stride_page_table); - auto mPT = mPT_l(_, batch_coord); - - int k_tile_total = ceil_div(K, TileShapeS{}); - int k_tile_per_cta = ceil_div(k_tile_total, split_kv); - int k_index = get<3>(blk_coord) * k_tile_per_cta; // lower limit - int k_tile_count = max(0, min(k_tile_total, k_index + k_tile_per_cta) - k_index); - if (k_tile_count == 0) { - return; - } - - auto page_size = Pow2{mainloop_args.page_size}; - auto pages_per_tile = Pow2{TileShapeS{} / page_size}; - int thread_idx = threadIdx.x % cutlass::NumThreadsPerWarp; - -#if 1 - for (; k_tile_count > 0; ++k_index, --k_tile_count) { - pipeline_page_table.producer_acquire(pipeline_pt_producer_state); - - // assume a single warp - - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < TileShapeS{}; i += cutlass::NumThreadsPerWarp) { - int idx = i + thread_idx; - bool guard = idx < pages_per_tile; - int smem_idx = pipeline_pt_producer_state.index() * TileShapeS::value + idx; - int pt_idx = pages_per_tile * k_index + idx; - - cutlass::arch::cp_async_zfill( - &shared_tensors.smem_page_table[smem_idx], &mPT(pt_idx), guard - ); - } - - pipeline_page_table.producer_commit(pipeline_pt_producer_state, cutlass::arch::cpasync_barrier_arrive); - ++pipeline_pt_producer_state; - } -#endif - } - - - struct Gather { - int& page_table_stage; - Pow2 pages_per_tile; - const int * __restrict__ smem_page_table; - - CUTLASS_DEVICE int operator()(int idx) const { - return smem_page_table[page_table_stage * TileShapeS::value + idx % pages_per_tile]; - } - - CUTLASS_DEVICE friend void print(Gather const&) { - printf(""); - } - - }; - - - template - CUTLASS_DEVICE void load_cpasync( - BlkCoord const& blk_coord, - ProblemShape const& problem_shape, - MainloopArguments const& mainloop_args, - MainloopParams const& mainloop_params, - TensorStorage& shared_tensors, - PipelineLoadQK& pipeline_load, - typename PipelineLoadQK::PipelineState& pipeline_load_producer_state, - int const& split_kv, - PipelinePT& pipeline_page_table, - typename PipelinePT::PipelineState& pipeline_pt_consumer_state) { - - auto [H, K, D, B] = problem_shape; - auto [D_latent, D_rope] = D; - - using X = Underscore; - - int k_tile_total = ceil_div(K, TileShapeS{}); - int k_tile_per_cta = ceil_div(k_tile_total, split_kv); - int k_index = get<3>(blk_coord) * k_tile_per_cta; // lower limit - int k_tile_count = max(0, min(k_tile_total, k_index + k_tile_per_cta) - k_index); - if (k_tile_count == 0) { - return; - } - - // partition all tensors - auto mQL = make_tensor(make_gmem_ptr(mainloop_args.ptr_q_latent), make_shape(H, D_latent, B), mainloop_args.stride_q_latent); - auto mQR = make_tensor(make_gmem_ptr(mainloop_args.ptr_q_rope), make_shape(H, D_rope, B), mainloop_args.stride_q_rope); - - int paged_B = mainloop_args.page_count; - auto paged_K = Pow2{mainloop_args.page_size}; - auto mPT_l = make_tensor(make_gmem_ptr(mainloop_args.ptr_page_table), make_shape(paged_B, B), mainloop_args.stride_page_table); - - int batch_coord = get<2>(blk_coord); - auto mPT = mPT_l(_, batch_coord); - - auto gQL = local_tile(mQL, TileShapeQK{}, make_coord(_,_,_), Step<_1, X, _1>{}); - auto gQR = local_tile(mQR, TileShapeQK{}, make_coord(_,_,_), Step<_1, X, _1>{}); - - ThrMMA cta_mma_qk = TiledMmaQK{}.get_slice(get<0>(blk_coord) % size(AtomThrShapeMNK{})); - ThrMMA cta_mma_pv = TiledMmaPV{}.get_slice(get<0>(blk_coord) % size(AtomThrShapeMNK{})); - - auto tSgQL = cta_mma_qk.partition_A(gQL); - auto tSgQR = cta_mma_qk.partition_A(gQR); - - Tensor sQ = make_tensor(make_smem_ptr(shared_tensors.smem_q.begin()), SmemLayoutQ{}); - Tensor sKC = make_tensor(make_smem_ptr(shared_tensors.smem_kc.begin()), SmemLayoutKC{}); - Tensor sVC = make_tensor(make_smem_ptr(shared_tensors.smem_vc.begin()), SmemLayoutVC{}); - - auto make_copy_for = [](auto sT) { - auto rT_a = sT.layout()(_, _, _, _0{}); - auto rT = make_ordered_layout(shape(rT_a), stride(rT_a)); - auto threads = Int{}; - auto values = Int{}; - return make_cotiled_copy( - Copy_Atom, Element>{}, - make_ordered_layout( - make_shape(threads, values), - make_stride(_1{}, _0{})), - rT); - }; - - // like cute::copy, but makes sure we do all page table lookups first - auto copy_split = [](auto atom, auto src, auto dst) { - auto src_v = group_modes<1, rank_v>(src); - auto dst_v = group_modes<1, rank_v>(dst); - - auto src_v_ptrs = make_tensor(size<1>(src_v)); - for (int i = 0; i < size<1>(src_v); i++) { - src_v_ptrs(i) = &src_v(_0{}, i); - } - - - for (int i = 0; i < size<1>(src_v); i++) { - auto src_v_i = make_tensor( - make_gmem_ptr(src_v_ptrs(i)), - make_shape(shape<0>(src_v)), - make_stride(make_stride(_1{}, _0{})) - ); - atom.call(src_v_i, dst_v(_, i)); - } - }; - - auto tiled_copy_q = make_copy_for(sQ); - auto tiled_copy_kc = make_copy_for(sKC); - auto tiled_copy_vc = make_copy_for(sVC); - - auto thr_copy_q = tiled_copy_q.get_thread_slice(threadIdx.x % (kNumLoadWarps * cutlass::NumThreadsPerWarp)); - auto thr_copy_kc = tiled_copy_kc.get_thread_slice(threadIdx.x % (kNumLoadWarps * cutlass::NumThreadsPerWarp)); - auto thr_copy_vc = tiled_copy_vc.get_thread_slice(threadIdx.x % (kNumLoadWarps * cutlass::NumThreadsPerWarp)); - - auto tQsQ = thr_copy_q.partition_D(sQ); - auto tQgQL = thr_copy_q.partition_S(tSgQL); - auto tQgQR = thr_copy_q.partition_S(tSgQR); - - auto tKCsKC = thr_copy_kc.partition_D(sKC); - auto tVCsVC = thr_copy_vc.partition_D(sVC); - - auto pipeline_pt_release_state = pipeline_pt_consumer_state; - - int page_table_stage = -1; - Pow2 pages_per_tile{TileShapeS{} / paged_K}; - const int * __restrict__ smem_page_table = shared_tensors.smem_page_table.begin(); - Gather gather{page_table_stage, pages_per_tile, smem_page_table}; - - auto mCL = make_tensor( - make_gmem_ptr(mainloop_args.ptr_c_latent), - ComposedLayout{ - make_layout( - make_shape(make_shape(paged_K, paged_B), _1{}), - make_stride(make_stride(get<0>(mainloop_args.stride_c_latent), example::CustomStride(gather, get<2>(mainloop_args.stride_c_latent))), get<1>(mainloop_args.stride_c_latent))), - make_coord(_0{}, _0{}), - make_identity_layout(make_shape(paged_K * paged_B, D_latent))}); - - auto mKR = make_tensor( - make_gmem_ptr(mainloop_args.ptr_k_rope), - ComposedLayout{ - make_layout( - make_shape(make_shape(paged_K, paged_B), _1{}), - make_stride(make_stride(get<0>(mainloop_args.stride_k_rope), example::CustomStride(gather, get<2>(mainloop_args.stride_k_rope))), get<1>(mainloop_args.stride_k_rope))), - make_coord(_0{}, _0{}), - make_identity_layout(make_shape(paged_K * paged_B, D_latent))}); - - auto mCLT = make_tensor( - make_gmem_ptr(mainloop_args.ptr_c_latent), - ComposedLayout{ - make_layout( - make_shape(_1{}, make_shape(paged_K, paged_B)), - make_stride(get<1>(mainloop_args.stride_c_latent), make_stride(get<0>(mainloop_args.stride_c_latent), example::CustomStride(gather, get<2>(mainloop_args.stride_c_latent))))), - make_coord(_0{}, _0{}), - make_identity_layout(make_shape(D_latent, paged_K * paged_B))}); - - auto gCL = local_tile(mCL, TileShapeQK{}, make_coord(_,_,_), Step{}); - auto gKR = local_tile(mKR, TileShapeQK{}, make_coord(_,_,_), Step{}); - auto gCLT = local_tile(mCLT, TileShapePV{}, make_coord(_,_,_), Step{}); - - auto tSgCL = cta_mma_qk.partition_B(gCL); - auto tSgKR = cta_mma_qk.partition_B(gKR); - auto tOgCLT = cta_mma_pv.partition_B(gCLT); - - auto tKCgCL = thr_copy_kc.partition_S(tSgCL); - auto tKCgKR = thr_copy_kc.partition_S(tSgKR); - auto tVCgCLT = thr_copy_vc.partition_S(tOgCLT); - - // latent is first in memory, so let's load it first always - // startup: alternate Q and K, set tx count appropriately, for k_idx = 0 - auto& pipeline_acquire_state = pipeline_load_producer_state; - auto pipeline_commit_state = pipeline_acquire_state; - int pipeline_offset = 0; - - for (int i = 0; i < StagesPV; i++) { - cutlass::arch::cp_async_fence(); - } - - auto load_stage = [&](auto fn) { - pipeline_load.producer_acquire(pipeline_acquire_state); - fn(pipeline_acquire_state.index()); - cutlass::arch::cp_async_fence(); - - ++pipeline_acquire_state; - ++pipeline_offset; - - if (pipeline_offset == StagesPV - 1) { - cutlass::arch::cp_async_wait(); - pipeline_load.producer_commit(pipeline_commit_state); - ++pipeline_commit_state; - --pipeline_offset; - } - }; - - pipeline_page_table.consumer_wait(pipeline_pt_consumer_state); - page_table_stage = pipeline_pt_consumer_state.index(); - ++pipeline_pt_consumer_state; - - // each Q/K tile consists of rope and latent - for (int i = 0; i < IterationsQKLatent; i++) { - load_stage([&](int index) { - cute::copy(tiled_copy_q, tQgQL(_, _, _, _, _0{}, i, batch_coord), tQsQ(_, _, _, _, i)); - copy_split(tiled_copy_kc, tKCgCL(_, _, _, _, k_index, i), tKCsKC(_, _, _, _, index)); - }); - } - - for (int i = 0; i < IterationsQKRope; i++) { - load_stage([&](int index) { - cute::copy(tiled_copy_q, tQgQR(_, _, _, _, _0{}, i, batch_coord), tQsQ(_, _, _, _, IterationsQKLatent + i)); - copy_split(tiled_copy_kc, tKCgKR(_, _, _, _, k_index, i), tKCsKC(_, _, _, _, index)); - }); - } - - k_index += 1; - k_tile_count -= 1; - - // assume k_tile_count >= 1 - // perform K+Q load here - CUTLASS_PRAGMA_NO_UNROLL - while (k_tile_count > 0) { - - pipeline_page_table.consumer_wait(pipeline_pt_consumer_state); - page_table_stage = pipeline_pt_consumer_state.index(); - ++pipeline_pt_consumer_state; - - for (int i = 0; i < IterationsQKLatent; i++) { - load_stage([&](int index) { - copy_split(tiled_copy_kc, tKCgCL(_, _, _, _, k_index, i), tKCsKC(_, _, _, _, index)); - }); - } - - for (int i = 0; i < IterationsQKRope; i++) { - load_stage([&](int index) { - copy_split(tiled_copy_kc, tKCgKR(_, _, _, _, k_index, i), tKCsKC(_, _, _, _, index)); - }); - } - - page_table_stage = pipeline_pt_release_state.index(); - - for (int i = 0; i < IterationsPV_K; i++) { - for (int j = 0; j < IterationsPV_N; j++) { - load_stage([&](int index) { - copy_split(tiled_copy_vc, tVCgCLT(_, _, _, _, j, IterationsPV_K * (k_index - 1) + i), tVCsVC(_, _, _, _, index)); - }); - } - } - - pipeline_page_table.consumer_release(pipeline_pt_release_state); - ++pipeline_pt_release_state; - - k_index += 1; - k_tile_count -= 1; - } - - page_table_stage = pipeline_pt_release_state.index(); - - for (int i = 0; i < IterationsPV_K; i++) { - for (int j = 0; j < IterationsPV_N; j++) { - load_stage([&](int index) { - copy_split(tiled_copy_vc, tVCgCLT(_, _, _, _, j, IterationsPV_K * (k_index - 1) + i), tVCsVC(_, _, _, _, index)); - }); - } - } - - pipeline_page_table.consumer_release(pipeline_pt_release_state); - ++pipeline_pt_release_state; - - while (pipeline_offset > 0) { - cutlass::arch::cp_async_fence(); - - cutlass::arch::cp_async_wait(); - pipeline_load.producer_commit(pipeline_commit_state); - ++pipeline_commit_state; - --pipeline_offset; - } - - cutlass::arch::cp_async_wait<0>(); - - } - - - template - CUTLASS_DEVICE void load_tma( - BlkCoord const& blk_coord, - ProblemShape const& problem_shape, - MainloopArguments const& mainloop_args, - MainloopParams const& mainloop_params, - TensorStorage& shared_tensors, - PipelineLoadQK& pipeline_load_qk, - typename PipelineLoadQK::PipelineState& pipeline_load_qk_producer_state, - PipelineLoadPV& pipeline_load_pv, - typename PipelineLoadPV::PipelineState& pipeline_load_pv_producer_state, - int const& split_kv) { - - auto [H, K, D, B] = problem_shape; - auto [D_latent, D_rope] = D; - - int k_tile_total = ceil_div(K, TileShapeS{}); - int k_tile_per_cta = ceil_div(k_tile_total, split_kv); - int k_index = get<3>(blk_coord) * k_tile_per_cta; // lower limit - int k_tile_count = max(0, min(k_tile_total, k_index + k_tile_per_cta) - k_index); - if (k_tile_count == 0) { - return; - } - - using X = Underscore; - - // partition all tensors - auto mQL = mainloop_params.tma_load_q_latent.get_tma_tensor(make_shape(H, D_latent, B)); - auto mQR = mainloop_params.tma_load_q_rope.get_tma_tensor(make_shape(H, D_rope, B)); - - int paged_B = B; - int paged_K = K; - if constexpr (kIsPaged) { - paged_B = mainloop_args.page_count; - paged_K = mainloop_args.page_size; - } - auto mPT_l = make_tensor(make_gmem_ptr(mainloop_args.ptr_page_table), make_shape(paged_B, B), mainloop_args.stride_page_table); - - auto mCL = mainloop_params.tma_load_c_latent.get_tma_tensor(make_shape(paged_K, D_latent, paged_B)); - auto mKR = mainloop_params.tma_load_k_rope.get_tma_tensor(make_shape(paged_K, D_rope, paged_B)); - - auto mCLT = mainloop_params.tma_load_c_latent_transpose.get_tma_tensor(make_shape(D_latent, paged_K, paged_B)); - - auto gQL = local_tile(mQL, TileShapeQK{}, make_coord(_,_,_), Step<_1, X, _1>{}); - auto gQR = local_tile(mQR, TileShapeQK{}, make_coord(_,_,_), Step<_1, X, _1>{}); - - auto gCL = local_tile(mCL, TileShapeQK{}, make_coord(_,_,_), Step{}); - auto gKR = local_tile(mKR, TileShapeQK{}, make_coord(_,_,_), Step{}); - auto gCLT = local_tile(mCLT, TileShapePV{}, make_coord(_,_,_), Step{}); - - ThrMMA cta_mma_qk = TiledMmaQK{}.get_slice(get<0>(blk_coord) % size(AtomThrShapeMNK{})); - ThrMMA cta_mma_pv = TiledMmaPV{}.get_slice(get<0>(blk_coord) % size(AtomThrShapeMNK{})); - - auto tSgQL = cta_mma_qk.partition_A(gQL); - auto tSgQR = cta_mma_qk.partition_A(gQR); - - auto tSgCL = cta_mma_qk.partition_B(gCL); - auto tSgKR = cta_mma_qk.partition_B(gKR); - - auto tOgCLT = cta_mma_pv.partition_B(gCLT); - - Tensor sQ = make_tensor(make_smem_ptr(shared_tensors.smem_q.begin()), SmemLayoutQ{}); - Tensor sKC = make_tensor(make_smem_ptr(shared_tensors.smem_kc.begin()), SmemLayoutKC{}); - Tensor sVC = make_tensor(make_smem_ptr(shared_tensors.smem_vc.begin()), SmemLayoutVC{}); - - auto [tQLgQL_mkl, tQsQ] = tma_partition( - mainloop_params.tma_load_q_latent, _0{}, make_layout(_1{}), - group_modes<0,3>(sQ), group_modes<0,3>(tSgQL)); - - auto [tQRgQR_mkl, tQsQ_ignore] = tma_partition( - mainloop_params.tma_load_q_rope, _0{}, make_layout(_1{}), - group_modes<0,3>(sQ), group_modes<0,3>(tSgQR)); - - auto [tCLgCL_nkl, tKCsKC] = tma_partition( - mainloop_params.tma_load_c_latent, _0{}, make_layout(_1{}), - group_modes<0,3>(sKC), group_modes<0,3>(tSgCL)); - - auto [tKRgKR_nkl, tKCsKC_ignore] = tma_partition( - mainloop_params.tma_load_k_rope, _0{}, make_layout(_1{}), - group_modes<0,3>(sKC), group_modes<0,3>(tSgKR)); - - auto [tCLTgCLT_nkl, tVCsVC] = tma_partition( - mainloop_params.tma_load_c_latent_transpose, _0{}, make_layout(_1{}), - group_modes<0,3>(sVC), group_modes<0,3>(tOgCLT)); - - uint16_t mcast_mask = 0; - - int batch_coord = get<2>(blk_coord); - Tensor tQLgQL = tQLgQL_mkl(_, _, _, batch_coord); - Tensor tQRgQR = tQRgQR_mkl(_, _, _, batch_coord); - - auto mPT = mPT_l(_, batch_coord); - - Tensor tCLgCL = tCLgCL_nkl(_, _, _, _); - Tensor tKRgKR = tKRgKR_nkl(_, _, _, _); - - // careful: stage and k are swapped here! - Tensor tCLTgCLT = tCLTgCLT_nkl(_, _, _, _); - - // latent is first in memory, so let's load it first always - // startup: alternate Q and K, set tx count appropriately, for k_idx = 0 - - // each Q/K tile consists of rope and latent - for (int i = 0; i < IterationsQKLatent; i++) { - pipeline_load_qk.producer_expect_transaction(pipeline_load_qk_producer_state, kTransactionsBytesLoadExtraQ); - pipeline_load_qk.producer_acquire(pipeline_load_qk_producer_state); - auto tma_barrier = pipeline_load_qk.producer_get_barrier(pipeline_load_qk_producer_state); - - if (cute::elect_one_sync()) { - // expect the extra bytes - // load_qk ql - cute::copy(mainloop_params.tma_load_q_latent.with(*tma_barrier, mcast_mask), tQLgQL(_, _0{}, i), tQsQ(_, i)); - // load_qk cl - if constexpr (kIsPaged) { - cute::copy( - mainloop_params.tma_load_c_latent.with(*tma_barrier, mcast_mask), - tCLgCL(_, _0{}, i, mPT(k_index)), - tKCsKC(_, pipeline_load_qk_producer_state.index()) - ); - } - else { - cute::copy( - mainloop_params.tma_load_c_latent.with(*tma_barrier, mcast_mask), - tCLgCL(_, k_index, i, batch_coord), - tKCsKC(_, pipeline_load_qk_producer_state.index())); - } - } - ++pipeline_load_qk_producer_state; - } - - for (int i = 0; i < IterationsQKRope; i++) { - pipeline_load_qk.producer_expect_transaction(pipeline_load_qk_producer_state, kTransactionsBytesLoadExtraQ); - pipeline_load_qk.producer_acquire(pipeline_load_qk_producer_state); - auto tma_barrier = pipeline_load_qk.producer_get_barrier(pipeline_load_qk_producer_state); - - if (cute::elect_one_sync()) { - // expect the extra bytes - // load_qk ql - cute::copy(mainloop_params.tma_load_q_rope.with(*tma_barrier, mcast_mask), tQRgQR(_, _0{}, i), tQsQ(_, i + IterationsQKLatent)); - // load_qk cl - if constexpr (kIsPaged) { - cute::copy( - mainloop_params.tma_load_k_rope.with(*tma_barrier, mcast_mask), - tKRgKR(_, _0{}, i, mPT(k_index)), - tKCsKC(_, pipeline_load_qk_producer_state.index()) - ); - } - else { - cute::copy( - mainloop_params.tma_load_k_rope.with(*tma_barrier, mcast_mask), - tKRgKR(_, k_index, i, batch_coord), - tKCsKC(_, pipeline_load_qk_producer_state.index())); - } - } - ++pipeline_load_qk_producer_state; - } - - k_index += 1; - k_tile_count -= 1; - - // assume k_tile_count >= 1 - // perform K+Q load here - CUTLASS_PRAGMA_NO_UNROLL - while (k_tile_count > 0) { - - // perform K load - for (int i = 0; i < IterationsQKLatent; i++) { - pipeline_load_qk.producer_acquire(pipeline_load_qk_producer_state); - auto tma_barrier = pipeline_load_qk.producer_get_barrier(pipeline_load_qk_producer_state); - - if (cute::elect_one_sync()) { - // load_qk cl - if constexpr (kIsPaged) { - cute::copy( - mainloop_params.tma_load_c_latent.with(*tma_barrier, mcast_mask), - tCLgCL(_, _0{}, i, mPT(k_index)), - tKCsKC(_, pipeline_load_qk_producer_state.index()) - ); - } - else { - cute::copy( - mainloop_params.tma_load_c_latent.with(*tma_barrier, mcast_mask), - tCLgCL(_, k_index, i, batch_coord), - tKCsKC(_, pipeline_load_qk_producer_state.index())); - } - } - ++pipeline_load_qk_producer_state; - } - - for (int i = 0; i < IterationsQKRope; i++) { - pipeline_load_qk.producer_acquire(pipeline_load_qk_producer_state); - auto tma_barrier = pipeline_load_qk.producer_get_barrier(pipeline_load_qk_producer_state); - - if (cute::elect_one_sync()) { - // load_qk cl - if constexpr (kIsPaged) { - cute::copy( - mainloop_params.tma_load_k_rope.with(*tma_barrier, mcast_mask), - tKRgKR(_, _0{}, i, mPT(k_index)), - tKCsKC(_, pipeline_load_qk_producer_state.index()) - ); - } - else { - cute::copy( - mainloop_params.tma_load_k_rope.with(*tma_barrier, mcast_mask), - tKRgKR(_, k_index, i, batch_coord), - tKCsKC(_, pipeline_load_qk_producer_state.index())); - } - } - ++pipeline_load_qk_producer_state; - } - - // prefetch next K load to keep busy while we transpose-load from cache - const int kPrefetchDistance = 1; - for (int i = 0; i < IterationsQKLatent; i++) { - if (cute::elect_one_sync()) { - if constexpr (kIsPaged) { - if (k_tile_count > kPrefetchDistance) { - cute::prefetch( - mainloop_params.tma_load_c_latent, - tCLgCL(_, _0{}, i, mPT(k_index + kPrefetchDistance)) - ); - } - } - else { - cute::prefetch( - mainloop_params.tma_load_c_latent, - tCLgCL(_, k_index + kPrefetchDistance, i, batch_coord) - ); - } - } - } - - for (int i = 0; i < IterationsQKRope; i++) { - if (cute::elect_one_sync()) { - if constexpr (kIsPaged) { - if (k_tile_count > kPrefetchDistance) { - cute::prefetch( - mainloop_params.tma_load_k_rope, - tKRgKR(_, _0{}, i, mPT(k_index + kPrefetchDistance)) - ); - } - } - else { - cute::prefetch( - mainloop_params.tma_load_k_rope, - tKRgKR(_, k_index + kPrefetchDistance, i, batch_coord) - ); - } - } - } - - // perform V load (k_idx - 1) - - for (int i = 0; i < IterationsPV_K; i++) { - for (int j = 0; j < IterationsPV_N; j++) { - pipeline_load_pv.producer_acquire(pipeline_load_pv_producer_state); - auto tma_barrier = pipeline_load_pv.producer_get_barrier(pipeline_load_pv_producer_state); - - if (cute::elect_one_sync()) { - // load_pv cl - // note the transpose in indices! - // note we are off-by-one on k_index - if constexpr (kIsPaged) { - cute::copy( - mainloop_params.tma_load_c_latent_transpose.with(*tma_barrier, mcast_mask, cute::TMA::CacheHintSm100::EVICT_FIRST), - tCLTgCLT(_, j, i, mPT(k_index - 1)), - tVCsVC(_, pipeline_load_pv_producer_state.index()) - ); - } - else { - cute::copy( - mainloop_params.tma_load_c_latent_transpose.with(*tma_barrier, mcast_mask, cute::TMA::CacheHintSm100::EVICT_FIRST), - tCLTgCLT(_, j, IterationsPV_K * (k_index - 1) + i, batch_coord), - tVCsVC(_, pipeline_load_pv_producer_state.index()) - ); - } - } - ++pipeline_load_pv_producer_state; - } - } - - k_index += 1; - k_tile_count -= 1; - } - - for (int i = 0; i < IterationsPV_K; i++) { - for (int j = 0; j < IterationsPV_N; j++) { - pipeline_load_pv.producer_acquire(pipeline_load_pv_producer_state); - auto tma_barrier = pipeline_load_pv.producer_get_barrier(pipeline_load_pv_producer_state); - - if (cute::elect_one_sync()) { - // load_pv cl - // note the transpose in indices - // note we are off-by-one on k_index - - if constexpr (kIsPaged) { - cute::copy( - mainloop_params.tma_load_c_latent_transpose.with(*tma_barrier, mcast_mask, cute::TMA::CacheHintSm100::EVICT_FIRST), - tCLTgCLT(_, j, i, mPT(k_index - 1)), - tVCsVC(_, pipeline_load_pv_producer_state.index()) - ); - } - else { - cute::copy( - mainloop_params.tma_load_c_latent_transpose.with(*tma_barrier, mcast_mask, cute::TMA::CacheHintSm100::EVICT_FIRST), - tCLTgCLT(_, j, IterationsPV_K * (k_index - 1) + i, batch_coord), - tVCsVC(_, pipeline_load_pv_producer_state.index()) - ); - } - } - ++pipeline_load_pv_producer_state; - } - } - } - - template - CUTLASS_DEVICE void mma( - BlkCoord const& blk_coord, - ProblemShape const& problem_shape, - TensorStorage& shared_tensors, - PipelineLoadQK& pipeline_load_qk, - typename PipelineLoadQK::PipelineState& pipeline_load_qk_consumer_state, - PipelineLoadPV& pipeline_load_pv, - typename PipelineLoadPV::PipelineState& pipeline_load_pv_consumer_state, - PipelineS& pipeline_mma_s, - typename PipelineS::PipelineState& pipeline_mma_s_producer_state, - PipelineP& pipeline_p_mma, - typename PipelineP::PipelineState& pipeline_p_mma_consumer_state, - PipelineO& pipeline_mma_o, - typename PipelineO::PipelineState& pipeline_mma_o_producer_state, - int const& split_kv) { - - auto [H, K, D, B] = problem_shape; - - int k_tile_total = ceil_div(K, TileShapeS{}); - int k_tile_per_cta = ceil_div(k_tile_total, split_kv); - int k_index = get<3>(blk_coord) * k_tile_per_cta; // lower limit - int k_tile_count = max(0, min(k_tile_total, k_index + k_tile_per_cta) - k_index); - if (k_tile_count == 0) { - return; - } - - // mma init - Tensor sQ = make_tensor(make_smem_ptr(shared_tensors.smem_q.begin()), SmemLayoutQ{}); - Tensor sKC = make_tensor(make_smem_ptr(shared_tensors.smem_kc.begin()), SmemLayoutKC{}); - Tensor sVC = make_tensor(make_smem_ptr(shared_tensors.smem_vc.begin()), SmemLayoutVC{}); - Tensor sP = make_tensor(make_smem_ptr((Element*) shared_tensors.smem_p.begin()), SmemLayoutP{}); - - Tensor tSrQ = TiledMmaQK::make_fragment_A(sQ); - Tensor tSrKC = TiledMmaQK::make_fragment_B(sKC); - Tensor tOrP = TiledMmaPV::make_fragment_A(sP); - Tensor tOrVC = TiledMmaPV::make_fragment_B(sVC); - - TiledMmaQK tiled_mma_qk; - TiledMmaPV tiled_mma_pv; - - Tensor tStS = partition_fragment_C(tiled_mma_qk, select<0,1>(TileShapeQK{})); - Tensor tOtO = partition_fragment_C(tiled_mma_pv, select<0,1>(TileShapePV{})); - - tiled_mma_pv.accumulate_ = UMMA::ScaleOut::Zero; - - pipeline_mma_s.producer_acquire(pipeline_mma_s_producer_state); - - // Mma S0 S1 O0 S2 O1 ... Sn On-1 On - // S0 ownership -- ----- -- -- - // S1 ownership -- ----- ---- - // O ownership -- -- ---- -- - - tiled_mma_qk.accumulate_ = UMMA::ScaleOut::Zero; - for (int i = 0; i < IterationsQK; i++) { - pipeline_load_qk.consumer_wait(pipeline_load_qk_consumer_state); - int read_stage = pipeline_load_qk_consumer_state.index(); - - tStS.data() = uint32_t(pipeline_mma_s_producer_state.index() == 0 ? TmemAllocation::kS0 : TmemAllocation::kS1); - - CUTLASS_PRAGMA_UNROLL - for (int k_block = 0; k_block < size<2>(tSrQ); ++k_block) { - cute::gemm(tiled_mma_qk, - tSrQ(_,_,k_block,i), - tSrKC(_,_,k_block,read_stage), - tStS); - tiled_mma_qk.accumulate_ = UMMA::ScaleOut::One; - } - - pipeline_load_qk.consumer_release(pipeline_load_qk_consumer_state); - ++pipeline_load_qk_consumer_state; - } - - pipeline_mma_s.producer_commit(pipeline_mma_s_producer_state); - ++pipeline_mma_s_producer_state; - - k_tile_count -= 1; - - CUTLASS_PRAGMA_NO_UNROLL - while (k_tile_count > 0) { - - pipeline_mma_s.producer_acquire(pipeline_mma_s_producer_state); - tiled_mma_qk.accumulate_ = UMMA::ScaleOut::Zero; - for (int i = 0; i < IterationsQK; i++) { - pipeline_load_qk.consumer_wait(pipeline_load_qk_consumer_state); - int read_stage = pipeline_load_qk_consumer_state.index(); - - tStS.data() = uint32_t(pipeline_mma_s_producer_state.index() == 0 ? TmemAllocation::kS0 : TmemAllocation::kS1); - - CUTLASS_PRAGMA_UNROLL - for (int k_block = 0; k_block < size<2>(tSrQ); ++k_block) { - cute::gemm(tiled_mma_qk, - tSrQ(_,_,k_block,i), - tSrKC(_,_,k_block,read_stage), - tStS); - tiled_mma_qk.accumulate_ = UMMA::ScaleOut::One; - } - - pipeline_load_qk.consumer_release(pipeline_load_qk_consumer_state); - ++pipeline_load_qk_consumer_state; - } - - pipeline_mma_s.producer_commit(pipeline_mma_s_producer_state); - ++pipeline_mma_s_producer_state; - - pipeline_mma_o.producer_acquire(pipeline_mma_o_producer_state); - pipeline_p_mma.consumer_wait(pipeline_p_mma_consumer_state); - - for (int i = 0; i < IterationsPV_K; i++) { - auto acc_flag = tiled_mma_pv.accumulate_; - for (int j = 0; j < IterationsPV_N; j++) { - pipeline_load_pv.consumer_wait(pipeline_load_pv_consumer_state); - - int read_stage = pipeline_load_pv_consumer_state.index(); - - tOtO.data() = uint32_t(TmemAllocation::kO0) + j * uint32_t(TmemAllocation::kSizeAccO); - tiled_mma_pv.accumulate_ = acc_flag; - - CUTLASS_PRAGMA_UNROLL - for (int k_block = 0; k_block < size<2>(tOrP); ++k_block) { - cute::gemm(tiled_mma_pv, - tOrP(_,_,k_block, make_coord(i, pipeline_p_mma_consumer_state.index())), - tOrVC(_,_,k_block,read_stage), - tOtO); - tiled_mma_pv.accumulate_ = UMMA::ScaleOut::One; - } - - pipeline_load_pv.consumer_release(pipeline_load_pv_consumer_state); - ++pipeline_load_pv_consumer_state; - } - } - - pipeline_p_mma.consumer_release(pipeline_p_mma_consumer_state); - ++pipeline_p_mma_consumer_state; - pipeline_mma_o.producer_commit(pipeline_mma_o_producer_state); - ++pipeline_mma_o_producer_state; - - --k_tile_count; - } - - pipeline_mma_o.producer_acquire(pipeline_mma_o_producer_state); - pipeline_p_mma.consumer_wait(pipeline_p_mma_consumer_state); - - for (int i = 0; i < IterationsPV_K; i++) { - auto acc_flag = tiled_mma_pv.accumulate_; - for (int j = 0; j < IterationsPV_N; j++) { - pipeline_load_pv.consumer_wait(pipeline_load_pv_consumer_state); - - int read_stage = pipeline_load_pv_consumer_state.index(); - - tOtO.data() = uint32_t(TmemAllocation::kO0) + j * uint32_t(TmemAllocation::kSizeAccO); - tiled_mma_pv.accumulate_ = acc_flag; - - CUTLASS_PRAGMA_UNROLL - for (int k_block = 0; k_block < size<2>(tOrP); ++k_block) { - cute::gemm(tiled_mma_pv, - tOrP(_,_,k_block, make_coord(i, pipeline_p_mma_consumer_state.index())), - tOrVC(_,_,k_block,read_stage), - tOtO); - tiled_mma_pv.accumulate_ = UMMA::ScaleOut::One; - } - - pipeline_load_pv.consumer_release(pipeline_load_pv_consumer_state); - ++pipeline_load_pv_consumer_state; - } - } - - pipeline_p_mma.consumer_release(pipeline_p_mma_consumer_state); - ++pipeline_p_mma_consumer_state; - pipeline_mma_o.producer_commit(pipeline_mma_o_producer_state); - ++pipeline_mma_o_producer_state; - } - - - template - CUTLASS_DEVICE void softmax( - IsLastTile const& is_last_tile, - ElementAcc& row_max, - ElementAcc& row_sum, - ElementAcc& correction_factor, - ProblemShape const& problem_shape, - MainloopArguments const& mainloop_args, - TensorStorage& shared_tensors, - int k_index, - uint32_t tmem_s, - int smem_p_index) { - - auto load_op = cute::SM100_TMEM_LOAD_32dp32b32x{}; - - TiledMmaQK tiled_mma_qk; - - Tensor tStS = partition_fragment_C(tiled_mma_qk, select<0,1>(TileShapeQK{})); - tStS.data() = tmem_s; - - CUTE_STATIC_ASSERT_V(shape<1>(tStS) == _1{}); - CUTE_STATIC_ASSERT_V(shape<2>(tStS) == _1{}); - Tensor tAcc = tStS(make_coord(_,_),_0{},_0{}); - - Tensor cS = make_identity_tensor(take<0,2>(CtaShapeQK{})); - - auto tiled_t2r = make_tmem_copy(load_op, tAcc); - auto thread_idx = threadIdx.x % size(tiled_t2r); - - auto thread_t2r = tiled_t2r.get_slice(thread_idx); - Tensor tTR_cS = thread_t2r.partition_D(cS); - Tensor tTR_rAcc = make_tensor(shape(tTR_cS)); - - Tensor tTR_rS_frag = make_tensor(shape(tTR_rAcc)); - const int AlignmentS = 4; - Tensor tTR_tAcc = thread_t2r.partition_S(tAcc); - Tensor tTR_rAcc_vec = recast>(tTR_rAcc); - Tensor tTR_rS_vec = recast>(tTR_rS_frag); - - // load s - copy(tiled_t2r, tTR_tAcc, tTR_rAcc); - - if (is_last_tile) { - for (int i = 0; i < size(tTR_rAcc); i++) { - if (get<1>(tTR_cS(i)) + TileShapeS{} * k_index >= get<1>(problem_shape)) { - tTR_rAcc(i) = -std::numeric_limits::infinity(); - } - } - } - - // max - ElementAcc row_max_new = row_max; - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(tTR_rAcc); i += 1) { - row_max_new = ::fmax(row_max_new, tTR_rAcc(i)); - } - - // for 2x2 dp, reduce here - if constexpr (kWarpsInN > 1) { - shared_tensors.smem_exchange[threadIdx.x] = row_max_new; - cutlass::arch::NamedBarrier(kNumComputeWarps*NumThreadsPerWarp, kNamedBarrierExchange).sync(); - // (64, 2) shape - int peer_index = (threadIdx.x + 64) % 128; - row_max_new = cutlass::max(row_max_new, shared_tensors.smem_exchange[peer_index]); - } - -#ifndef B2B - // find correction factor - ElementAcc softmax_scale_log2 = mainloop_args.softmax_scale * static_cast(M_LOG2E); - correction_factor = ::exp2f(softmax_scale_log2 * (row_max - row_max_new)); - row_max = row_max_new; - - // softmax - ElementAcc row_max_scale_log2 = row_max * softmax_scale_log2; - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(tTR_rAcc); i++) { - tTR_rAcc(i) = ::exp2f(softmax_scale_log2 * tTR_rAcc(i) - row_max_scale_log2); - } -#endif - - // quantize - cutlass::NumericArrayConverter epilogue_op; - - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(tTR_rAcc_vec); i++) { - tTR_rS_vec(i) = epilogue_op(tTR_rAcc_vec(i)); - } - - Tensor sP = make_tensor(make_smem_ptr((Element*) shared_tensors.smem_p.begin()), SmemLayoutP{})(_, _, _, make_coord(_, smem_p_index)); - - Tensor tOcP = TiledMmaPV{}.get_slice(_0{}).partition_A(cS); - - // have a mapping for each thread to coord - // find identical mapping to coords for the MMA - auto l = make_ordered_layout(make_shape(make_shape(_64{}, _2{}), make_shape(_16{}, TileShapeS{} / _32{})), make_stride(make_stride(_0{}, _3{}), make_stride(_1{}, _2{}))); - auto sP_ = as_position_independent_swizzle_tensor(sP); - copy_aligned(tTR_rS_frag, sP_.compose(l)(threadIdx.x, _)); - - // sum - row_sum *= correction_factor; - - static_assert(cute::is_same_v); - auto tTR_rAcc_float2 = recast(tTR_rAcc); - auto sums = make_tensor(_4{}); - static_assert(size(tTR_rAcc_float2) % size(sums) == 0); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(sums); i++) { - sums(i) = tTR_rAcc_float2(i); - } - CUTLASS_PRAGMA_UNROLL - for (int i = size(sums); i < size(tTR_rAcc_float2); i += size(sums)) { - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < size(sums); j++) { - cute::add(sums(j), sums(j), tTR_rAcc_float2(i + j)); - } - } - CUTLASS_PRAGMA_UNROLL - for (int i = 1; i < size(sums); i *= 2) { - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < size(sums); j += 2*i) { - cute::add(sums(j), sums(j), sums(j+i)); - } - } - row_sum += sums(0).x + sums(0).y; - } - - - CUTLASS_DEVICE void rescale( - ElementAcc correction_factor, - uint32_t tmem_o) { - - // for b2b gemm, do nothing -#ifndef B2B - auto load_op = cute::SM100_TMEM_LOAD_32dp32b32x{}; - auto store_op = TMEM::tmem_load_to_store(load_op); - - TiledMmaPV tiled_mma_pv; - - Tensor tOtO = partition_fragment_C(tiled_mma_pv, select<0,1>(TileShapePV{})); - tOtO.data() = tmem_o; - - CUTE_STATIC_ASSERT_V(shape<1>(tOtO) == _1{}); - CUTE_STATIC_ASSERT_V(shape<2>(tOtO) == _1{}); - Tensor tAcc = tOtO(make_coord(_,_),_0{},_0{}); - - auto cta_tiler_pv = take<0,2>(typename CollectiveMmaPV::CtaShape_MNK{}); - Tensor gO = make_tensor(make_gmem_ptr((ElementAcc*) nullptr), cta_tiler_pv, make_stride(0, 0)); - - auto tiled_t2r = make_tmem_copy(load_op, tAcc); - auto tiled_r2t = make_tmem_copy(store_op, tAcc); - auto thread_idx = threadIdx.x % size(tiled_t2r); - - auto thread_t2r = tiled_t2r.get_slice(thread_idx); - auto thread_r2t = tiled_r2t.get_slice(thread_idx); - Tensor tTR_gO = thread_t2r.partition_D(gO); - Tensor tTR_rAcc = make_tensor(shape(tTR_gO)); - - Tensor tTR_tAcc = thread_t2r.partition_S(tAcc); - - // load o - copy(tiled_t2r, tTR_tAcc, tTR_rAcc); - - // multiply by correction factor - float2 correction_factor_vec = make_float2(correction_factor, correction_factor); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(tTR_rAcc); i += 2) { - float2 in = make_float2(tTR_rAcc(i + 0), tTR_rAcc(i + 1)); - float2 out; - cute::mul(out, in, correction_factor_vec); - tTR_rAcc(i + 0) = out.x; - tTR_rAcc(i + 1) = out.y; - } - - // store o - copy(tiled_r2t, tTR_rAcc, tTR_tAcc); -#endif - } - - - template - CUTLASS_DEVICE void epilogue( - ElementAcc& row_max, - ElementAcc& row_sum, - BlkCoord const& cta_coord, - ProblemShape const& problem_shape, - MainloopArguments const& mainloop_args, - EpilogueParams const& epilogue_args, - TensorStorage& shared_tensors, - uint32_t tmem_o, - int const& split_kv) { - - auto load_op = cute::SM100_TMEM_LOAD_32dp32b32x{}; - - TiledMmaPV tiled_mma_pv; - - Tensor tOtO = TiledMmaPV::make_fragment_C(partition_shape_C(TiledMmaPV{}, take<0, 2>(TileShapePV{}))); - tOtO.data() = tmem_o; - - CUTE_STATIC_ASSERT_V(shape<1>(tOtO) == _1{}); - CUTE_STATIC_ASSERT_V(shape<2>(tOtO) == _1{}); - Tensor tAcc = tOtO(make_coord(_,_),_0{},_0{}); - - auto [H, K, D, B] = problem_shape; - auto [D_latent, D_rope] = D; - if (epilogue_args.ptr_o_acc != nullptr) { - using ElementOutAcc = ElementAcc; - constexpr auto AlignmentOutAcc = 128 / cute::sizeof_bits_v; - Tensor mO = make_tensor(make_gmem_ptr(epilogue_args.ptr_o_acc + get<3>(cta_coord) * D_latent), make_shape(H, D_latent, B), epilogue_args.stride_o_acc); - auto cta_tiler_pv = take<0,2>(typename CollectiveMmaPV::CtaShape_MNK{}); - Tensor gO = local_tile(mO, cta_tiler_pv, take<0,3>(cta_coord)); - - auto tiled_t2r = make_tmem_copy(load_op, tAcc); - auto thread_idx = threadIdx.x % size(tiled_t2r); - - auto thread_t2r = tiled_t2r.get_slice(thread_idx); - Tensor tTR_gO = thread_t2r.partition_D(gO); - Tensor tTR_rAcc = make_tensor(shape(tTR_gO)); - - Tensor tTR_rO_frag = make_tensor(shape(tTR_rAcc)); - Tensor tTR_rO_src = recast>(coalesce(tTR_rO_frag)); - Tensor tR2G_rO_dst = recast>(coalesce(tTR_gO)); - Tensor tTR_tAcc = thread_t2r.partition_S(tAcc); - - copy(tiled_t2r, tTR_tAcc, tTR_rAcc); - - cutlass::epilogue::thread::LinearCombination epilogue_op({epilogue_args.output_scale / row_sum}); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(tTR_rAcc); i++) { - tTR_rO_frag(i) = epilogue_op(tTR_rAcc(i)); - } - - copy(tTR_rO_src, tR2G_rO_dst); - -#ifndef B2B - - // compute LSE - ElementAcc lse = cutlass::fast_log(row_sum) + mainloop_args.softmax_scale * row_max; - - // store LSE - Tensor mLSE = make_tensor(make_gmem_ptr(epilogue_args.ptr_lse_acc + H * get<3>(cta_coord)), make_shape(H, B), epilogue_args.stride_lse_acc); - Tensor gLSE = local_tile(mLSE, append<3>(cta_tiler_pv, _1{}), take<0,3>(cta_coord), Step<_1, Underscore, _1>{}); - // for 2x2 dp, this must be conditional and the index is wrong - if (! kIs2Sm || (threadIdx.x < 64)) - { - gLSE(threadIdx.x) = lse; - } - #endif - } - else { - Tensor mO = make_tensor(make_gmem_ptr(epilogue_args.ptr_o), make_shape(H, D_latent, B), epilogue_args.stride_o); - auto cta_tiler_pv = take<0,2>(typename CollectiveMmaPV::CtaShape_MNK{}); - Tensor gO = local_tile(mO, cta_tiler_pv, take<0,3>(cta_coord)); - - auto tiled_t2r = make_tmem_copy(load_op, tAcc); - auto thread_idx = threadIdx.x % size(tiled_t2r); - - auto thread_t2r = tiled_t2r.get_slice(thread_idx); - Tensor tTR_gO = thread_t2r.partition_D(gO); - Tensor tTR_rAcc = make_tensor(shape(tTR_gO)); - - Tensor tTR_rO_frag = make_tensor(shape(tTR_rAcc)); - Tensor tTR_rO_src = recast>(coalesce(tTR_rO_frag)); - Tensor tR2G_rO_dst = recast>(coalesce(tTR_gO)); - Tensor tTR_tAcc = thread_t2r.partition_S(tAcc); - - copy(tiled_t2r, tTR_tAcc, tTR_rAcc); - - cutlass::epilogue::thread::LinearCombination epilogue_op({epilogue_args.output_scale / row_sum}); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(tTR_rAcc); i++) { - tTR_rO_frag(i) = epilogue_op(tTR_rAcc(i)); - } - - copy(tTR_rO_src, tR2G_rO_dst); - -#ifndef B2B - if (epilogue_args.ptr_lse != nullptr) { - // compute LSE - ElementAcc lse = cutlass::fast_log(row_sum) + mainloop_args.softmax_scale * row_max; - - // store LSE - Tensor mLSE = make_tensor(make_gmem_ptr(epilogue_args.ptr_lse), make_shape(H, B), epilogue_args.stride_lse); - Tensor gLSE = local_tile(mLSE, append<3>(cta_tiler_pv, _1{}), take<0,3>(cta_coord), Step<_1, Underscore, _1>{}); - - // for 2x2 dp, this must be conditional and the index is wrong - if (! kIs2Sm || (threadIdx.x < 64)) - { - gLSE(threadIdx.x) = lse; - } - } -#endif - } - } - - - template - CUTLASS_DEVICE void compute( - CtaCoord const& cta_coord, - ProblemShape const& problem_shape, - MainloopArguments const& mainloop_args, - EpilogueParams const& epilogue_args, - TensorStorage& shared_tensors, - PipelineS& pipeline_mma_s, - typename PipelineS::PipelineState& pipeline_mma_s_consumer_state, - PipelineP& pipeline_p_mma, - typename PipelineP::PipelineState& pipeline_p_mma_producer_state, - PipelineO& pipeline_mma_o, - typename PipelineO::PipelineState& pipeline_mma_o_consumer_state, - int const& split_kv) { - - auto [H, K, D, B] = problem_shape; - - int k_tile_total = ceil_div(K, TileShapeS{}); - int k_tile_per_cta = ceil_div(k_tile_total, split_kv); - int k_index = get<3>(cta_coord) * k_tile_per_cta; // lower limit - int k_tile_count = max(0, min(k_tile_total, k_index + k_tile_per_cta) - k_index); - if (k_tile_count == 0) { - - // if we return early, we have to make sure we release the load warp - cutlass::arch::NamedBarrier( - (kNumComputeWarps + kNumLoadWarps) * NumThreadsPerWarp, - kNamedBarrierEpilogue - ).arrive(); - - return; - } - int k_index_final = k_tile_total - 1; - - ElementAcc row_max = -std::numeric_limits::infinity(); - ElementAcc row_sum = 0; - ElementAcc correction_factor = 1; - - pipeline_p_mma.producer_acquire(pipeline_p_mma_producer_state); - pipeline_mma_s.consumer_wait(pipeline_mma_s_consumer_state); - - auto dispatch_bool = [](bool b, auto fn) { - if (b) { - fn(cute::true_type{}); - } - else { - fn(cute::false_type{}); - } - }; - - // softmax s0 -> p0 - dispatch_bool(k_index == k_index_final, [&](auto is_last_tile) { - softmax( - is_last_tile, - row_max, row_sum, correction_factor, - problem_shape, mainloop_args, shared_tensors, k_index, - uint32_t(pipeline_mma_s_consumer_state.index() == 0 ? TmemAllocation::kS0 : TmemAllocation::kS1), - pipeline_p_mma_producer_state.index() - ); - }); - - k_index += 1; - - cutlass::arch::fence_view_async_tmem_load(); - cutlass::arch::fence_view_async_shared(); - pipeline_mma_s.consumer_release(pipeline_mma_s_consumer_state); - ++pipeline_mma_s_consumer_state; - pipeline_p_mma.producer_commit(pipeline_p_mma_producer_state); - ++pipeline_p_mma_producer_state; - - k_tile_count -= 1; - - CUTLASS_PRAGMA_NO_UNROLL - while (k_tile_count > 0) { - pipeline_p_mma.producer_acquire(pipeline_p_mma_producer_state); - pipeline_mma_s.consumer_wait(pipeline_mma_s_consumer_state); - - // softmax s1 -> p1 - dispatch_bool(k_index == k_index_final, [&](auto is_last_tile) { - softmax( - is_last_tile, - row_max, row_sum, correction_factor, - problem_shape, mainloop_args, shared_tensors, k_index, - uint32_t(pipeline_mma_s_consumer_state.index() == 0 ? TmemAllocation::kS0 : TmemAllocation::kS1), - pipeline_p_mma_producer_state.index() - ); - }); - - cutlass::arch::fence_view_async_tmem_load(); - cutlass::arch::fence_view_async_shared(); - pipeline_mma_s.consumer_release(pipeline_mma_s_consumer_state); - ++pipeline_mma_s_consumer_state; - pipeline_p_mma.producer_commit(pipeline_p_mma_producer_state); - ++pipeline_p_mma_producer_state; - - pipeline_mma_o.consumer_wait(pipeline_mma_o_consumer_state); - - // rescale - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < IterationsPV_N; j++) { - rescale(correction_factor, uint32_t(TmemAllocation::kO0) + j * uint32_t(TmemAllocation::kSizeAccO)); - } - - cutlass::arch::fence_view_async_tmem_store(); - pipeline_mma_o.consumer_release(pipeline_mma_o_consumer_state); - ++pipeline_mma_o_consumer_state; - - --k_tile_count; - k_index += 1; - } - - pipeline_mma_o.consumer_wait(pipeline_mma_o_consumer_state); - -#ifdef B2B - row_sum = 1; -#else - if constexpr (kWarpsInN > 1) { - // reduce row_sum if needed (for 2x2 dp) - shared_tensors.smem_exchange[threadIdx.x] = row_sum; - cutlass::arch::NamedBarrier(kNumComputeWarps*NumThreadsPerWarp, kNamedBarrierExchange).sync(); - // (64, 2) shape - int peer_index = (threadIdx.x + 64) % 128; - row_sum += shared_tensors.smem_exchange[peer_index]; - } -#endif - - cutlass::arch::NamedBarrier((kNumComputeWarps + kNumLoadWarps) * NumThreadsPerWarp, kNamedBarrierEpilogue).arrive(); - - // epilogue - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < IterationsPV_N; j++) { - epilogue( - row_max, row_sum, - replace<1>(cta_coord, j), problem_shape, - mainloop_args, epilogue_args, shared_tensors, - uint32_t(TmemAllocation::kO0) + j * uint32_t(TmemAllocation::kSizeAccO), split_kv - ); - } - - cutlass::arch::fence_view_async_tmem_load(); - pipeline_mma_o.consumer_release(pipeline_mma_o_consumer_state); - ++pipeline_mma_o_consumer_state; - } - -}; - -/////////////////////////////////////////////////////////////////////////////// - -} // namespace cutlass::fmha::kernel diff --git a/python/sglang/kernels/aot/csrc/attention/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp b/python/sglang/kernels/aot/csrc/attention/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp deleted file mode 100644 index 30389e79f..000000000 --- a/python/sglang/kernels/aot/csrc/attention/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp +++ /dev/null @@ -1,160 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// clang-format off -#pragma once - -#include "cutlass/cutlass.h" -#include "cutlass/fast_math.h" -#include "cutlass/kernel_hardware_info.h" - -namespace cutlass::fmha::kernel { - -//////////////////////////////////////////////////////////////////////////////// - -struct Sm100MlaIndividualTileScheduler { - - struct Params { - dim3 grid; - }; - - bool valid_ = true; - - CUTLASS_DEVICE - Sm100MlaIndividualTileScheduler(Params const&) {} - - template - static Params to_underlying_arguments( - ProblemShape const& problem_shape, KernelHardwareInfo hw_info, - ClusterShape const& cluster_shape, int const& split_kv) { - using namespace cute; - dim3 grid(get<0>(cluster_shape), get<3>(problem_shape) /* Batch */, split_kv /*Maximum Split KV*/); - return Params{ grid }; - } - - static dim3 get_grid_shape(Params const& params) { - return params.grid; - } - - CUTLASS_DEVICE - bool is_valid() { - return valid_; - } - - CUTLASS_DEVICE - auto get_block_coord() { - using namespace cute; - return make_coord(blockIdx.x, _0{}, blockIdx.y, blockIdx.z); - } - - CUTLASS_DEVICE - Sm100MlaIndividualTileScheduler& operator++() { - valid_ = false; - return *this; - } -}; - -//////////////////////////////////////////////////////////////////////////////// - -struct Sm100MlaPersistentTileScheduler { - - struct Params { - int num_blocks; - FastDivmod divmod_m_block; - FastDivmod divmod_b; - FastDivmod divmod_split_kv; - KernelHardwareInfo hw_info; - }; - - int block_idx = 0; - Params params; - - CUTLASS_DEVICE - Sm100MlaPersistentTileScheduler(Params const& params) : block_idx(blockIdx.x), params(params) {} - - template - static Params to_underlying_arguments( - ProblemShape const& problem_shape, KernelHardwareInfo hw_info, - ClusterShape const& cluster_shape, int const& split_kv) { - using namespace cute; - // Get SM count if needed, otherwise use user supplied SM count - int sm_count = hw_info.sm_count; - if (sm_count <= 1 || sm_count % size<0>(cluster_shape) != 0) { - CUTLASS_TRACE_HOST(" WARNING: Arguments do not include a valid SM count.\n" - " For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count."); - sm_count = KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); - } - - CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count); - hw_info.sm_count = sm_count; - - int num_m_blocks = size<0>(cluster_shape); - int num_blocks = num_m_blocks * get<3>(problem_shape) /* Batch */; - num_blocks *= split_kv; /* Maximum Split KV*/ - - return Params { - num_blocks, - { num_m_blocks}, { get<3>(problem_shape) }, {split_kv}, - hw_info - }; - } - - static dim3 get_grid_shape(Params const& params) { - dim3 grid(std::min(params.num_blocks, params.hw_info.sm_count), 1, 1); - return grid; - } - - CUTLASS_DEVICE - bool is_valid() { - return block_idx < params.num_blocks; - } - - CUTLASS_DEVICE - auto get_block_coord() { - using namespace cute; - int block_decode = block_idx; - int m_block, bidb, n_split_kv; - params.divmod_m_block(block_decode, m_block, block_decode); - params.divmod_b(block_decode, bidb, block_decode); - params.divmod_split_kv(block_decode, n_split_kv, block_decode); - return make_coord(m_block, _0{}, bidb, n_split_kv); - } - - CUTLASS_DEVICE - Sm100MlaPersistentTileScheduler& operator++() { - block_idx += gridDim.x; - return *this; - } -}; - -//////////////////////////////////////////////////////////////////////////////// - -} // namespace cutlass::fmha::kernel diff --git a/python/sglang/kernels/aot/csrc/attention/vertical_slash_index.cu b/python/sglang/kernels/aot/csrc/attention/vertical_slash_index.cu deleted file mode 100644 index 118f780dd..000000000 --- a/python/sglang/kernels/aot/csrc/attention/vertical_slash_index.cu +++ /dev/null @@ -1,462 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT license. -// This file is for blocksparse attention utils cuda kernel. - -#include -#include -#include -#include - -// Save the start index of each block in the given range into block_offset. -// Returns the updated block count. -__device__ int64_t save_blocks( - int* block_offset, - int64_t range_start, - int64_t range_end, - int64_t block_size, - int64_t input_block_count, - int64_t kv_seqlen) { - if (range_start >= kv_seqlen) { - return input_block_count; - } - if (range_end > kv_seqlen) { - range_end = kv_seqlen; - } - int64_t current_block_count = input_block_count; - for (int idx = range_start; idx < range_end; idx += block_size) { - block_offset[current_block_count++] = idx; - } - return current_block_count; -} - -// CUDA kernel: convert sparse vertical/slash indices to block/column offsets. -__global__ void convert_vertical_slash_indexes_kernel( - const int* q_seqlens, // [BATCH, ] - const int* kv_seqlens, // [BATCH, ] - const int* vertical_indexes, // [BATCH, N_HEADS, NNZ_V] - const int* slash_indexes, // [BATCH, N_HEADS, NNZ_S] - int* block_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)] - int* block_offset, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S] - int* column_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)] - int* column_index, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V] - int64_t N_HEADS, - int64_t N_ROWS, - int64_t BLOCK_SIZE_M, - int64_t BLOCK_SIZE_N, - int64_t NNZ_V, - int64_t NNZ_S, - bool causal // True for intra, False for succ -) { - const int batch_idx = blockIdx.y; - const int head_idx = blockIdx.x; - const int group_idx = blockIdx.z; - - int64_t q_seqlen = q_seqlens[batch_idx]; - int64_t kv_seqlen = kv_seqlens[batch_idx]; - int64_t block_idx_m = group_idx * blockDim.x + threadIdx.x; - int64_t start_m = block_idx_m * BLOCK_SIZE_M; - if (start_m >= q_seqlen) { - return; - } - int64_t end_m = start_m + BLOCK_SIZE_M; - vertical_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_V; - slash_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_S; - int64_t row_offset = (batch_idx * N_HEADS + head_idx) * N_ROWS + block_idx_m; - block_count += row_offset; - block_offset += row_offset * NNZ_S; - column_count += row_offset; - column_index += row_offset * NNZ_V; - - bool has_slash = true; - int64_t tmp_col_cnt = 0, tmp_blk_cnt = 0; - int64_t s = 0, v = 0; - int64_t v_idx = vertical_indexes[v++]; - int64_t s_idx = slash_indexes[s++]; - if (causal) { - while (s_idx >= end_m + (kv_seqlen - q_seqlen) && s < NNZ_S) { - s_idx = slash_indexes[s++]; - } - if (s_idx > end_m + (kv_seqlen - q_seqlen)) has_slash = false; - s_idx = max((kv_seqlen - q_seqlen) + end_m - s_idx, BLOCK_SIZE_M); - } else { - while (s_idx >= end_m + kv_seqlen && s < NNZ_S) { - s_idx = slash_indexes[s++]; - } - if (s_idx > end_m + kv_seqlen) has_slash = false; - s_idx = max(kv_seqlen + end_m - s_idx, BLOCK_SIZE_M); - } - - int64_t range_start = s_idx - BLOCK_SIZE_M, range_end = s_idx; - if (!has_slash) { - if (causal) { - range_start = (kv_seqlen - q_seqlen) + end_m; - range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N; - } else { - range_start = kv_seqlen; - range_end = kv_seqlen + BLOCK_SIZE_N; - } - } - - bool slash_finished = false; - while (1) { - if (v_idx < range_end) { - if (v_idx < range_start) { - column_index[tmp_col_cnt++] = v_idx; - } - if (v < NNZ_V) { - v_idx = vertical_indexes[v++]; - } else { - if (causal) - v_idx = end_m + BLOCK_SIZE_N + (kv_seqlen - q_seqlen); - else - v_idx = end_m + BLOCK_SIZE_N + kv_seqlen; - } - } else { - if ((s < NNZ_S && causal) || (s < NNZ_S && !causal && slash_indexes[s] >= start_m)) { - if (causal) - s_idx = max((kv_seqlen - q_seqlen) + end_m - slash_indexes[s++], BLOCK_SIZE_M); - else - s_idx = max(kv_seqlen + end_m - slash_indexes[s++], BLOCK_SIZE_M); - } else { - if (v == NNZ_V || (v_idx > range_start && causal)) { - // add the last vertical if no more slash - if (v == NNZ_V && !causal && v_idx < kv_seqlen) { - column_index[tmp_col_cnt++] = v_idx; - } - tmp_blk_cnt = save_blocks(block_offset, range_start, range_end, BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen); - break; - } else { - if (causal) { - range_start = (kv_seqlen - q_seqlen) + end_m; - range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N; - } else { - // if slash_finished but there are vertical left, save current - // blocks - tmp_blk_cnt = save_blocks(block_offset, range_start, range_end, BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen); - range_start = kv_seqlen; - range_end = kv_seqlen + BLOCK_SIZE_N; - } - slash_finished = true; - } - } - if (!slash_finished) { - if (s_idx > range_end + BLOCK_SIZE_M) { - tmp_blk_cnt = save_blocks(block_offset, range_start, range_end, BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen); - range_start = s_idx - BLOCK_SIZE_M; - range_end = s_idx; - } else if (s_idx > range_end) { - range_end += BLOCK_SIZE_M; - } - } - } - } - - block_count[0] = tmp_blk_cnt; - column_count[0] = tmp_col_cnt; -} - -// Host function: launches the kernel with 64 threads per block. -void convert_vertical_slash_indexes_64x64( - const int* q_seqlens, // [BATCH, ] - const int* kv_seqlens, // [BATCH, ] - const int* vertical_indexes, // [BATCH, N_HEADS, NNZ_V] - const int* slash_indexes, // [BATCH, N_HEADS, NNZ_S] - int* block_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)] - int* block_offset, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S] - int* column_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)] - int* column_index, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V] - int64_t BATCH_SIZE, - int64_t N_HEADS, - int64_t N_ROWS, - int64_t BLOCK_SIZE_M, - int64_t BLOCK_SIZE_N, - int64_t NNZ_V, - int64_t NNZ_S, - bool causal) { - const int N_THREADS = 64; - const dim3 dimBlock((int32_t)N_THREADS); - const dim3 dimGrid( - (int32_t)N_HEADS, (int32_t)BATCH_SIZE, ((int32_t)N_ROWS + (int32_t)N_THREADS - 1) / (int32_t)N_THREADS); - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - convert_vertical_slash_indexes_kernel<<>>( - q_seqlens, - kv_seqlens, - vertical_indexes, - slash_indexes, - block_count, - block_offset, - column_count, - column_index, - N_HEADS, - N_ROWS, - BLOCK_SIZE_M, - BLOCK_SIZE_N, - NNZ_V, - NNZ_S, - causal); -} - -// Host function: prepares tensor pointers and launches the CUDA kernel. -void convert_vertical_slash_indexes( - torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS] - torch::Tensor& block_offset, // [BATCH, N_HEADS, NUM_ROWS, NNZ_S] - torch::Tensor& column_count, // [BATCH, N_HEADS, NUM_ROWS] - torch::Tensor& column_index, // [BATCH, N_HEADS, NUM_ROWS, NNZ_V] - torch::Tensor q_seqlens, // [BATCH, ] - torch::Tensor kv_seqlens, // [BATCH, ] - torch::Tensor vertical_indexes, // [BATCH, N_HEADS, NNZ_V] - torch::Tensor slash_indexes, // [BATCH, N_HEADS, NNZ_S] - int64_t context_size, - int64_t block_size_M, - int64_t block_size_N, - bool causal) { - cudaSetDevice(q_seqlens.get_device()); - - int64_t batch_size = slash_indexes.size(0); - int64_t num_heads = slash_indexes.size(1); - int64_t nnz_slash = slash_indexes.size(2); - int64_t nnz_vertical = vertical_indexes.size(2); - int64_t num_rows = (context_size + block_size_M - 1) / block_size_M; - - convert_vertical_slash_indexes_64x64( - q_seqlens.data_ptr(), - kv_seqlens.data_ptr(), - vertical_indexes.data_ptr(), - slash_indexes.data_ptr(), - block_count.data_ptr(), - block_offset.data_ptr(), - column_count.data_ptr(), - column_index.data_ptr(), - batch_size, - num_heads, - num_rows, - block_size_M, - block_size_N, - nnz_vertical, - nnz_slash, - causal); -} - -// --- mergehead kernels --- // - -// Kernel: like above, but supports per-head variable NNZ_V/NNZ_S. -__global__ void convert_vertical_slash_indexes_kernel_mergehead( - const int* q_seqlens, // [BATCH, ] - const int* kv_seqlens, // [BATCH, ] - const int* vertical_indexes, // [BATCH, N_HEADS, NNZ_V] - const int* slash_indexes, // [BATCH, N_HEADS, NNZ_S] - const int* per_head_vertical_topkv, - const int* per_head_slash_topkv, - int* block_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)] - int* block_offset, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S] - int* column_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)] - int* column_index, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V] - int64_t N_HEADS, - int64_t N_ROWS, - int64_t BLOCK_SIZE_M, - int64_t BLOCK_SIZE_N, - int64_t NNZ_V, - int64_t NNZ_S, - bool causal // True for intra, False for succ -) { - const int batch_idx = blockIdx.y; - const int head_idx = blockIdx.x; - const int group_idx = blockIdx.z; - - int64_t q_seqlen = q_seqlens[batch_idx]; - int64_t kv_seqlen = kv_seqlens[batch_idx]; - int64_t block_idx_m = group_idx * blockDim.x + threadIdx.x; - int64_t start_m = block_idx_m * BLOCK_SIZE_M; - if (start_m >= q_seqlen) { - return; - } - int64_t end_m = start_m + BLOCK_SIZE_M; - vertical_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_V; - slash_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_S; - int64_t row_offset = (batch_idx * N_HEADS + head_idx) * N_ROWS + block_idx_m; - block_count += row_offset; - block_offset += row_offset * NNZ_S; - column_count += row_offset; - column_index += row_offset * NNZ_V; - - // MergeHead: each head has it's unique max topk NNZ_V，NNZ_S. (NNZ_V，NNZ_S - // above is buffer size, use to compute offset) - NNZ_S = per_head_slash_topkv[head_idx]; - NNZ_V = per_head_vertical_topkv[head_idx]; - - bool has_slash = true; - int64_t tmp_col_cnt = 0, tmp_blk_cnt = 0; - int64_t s = 0, v = 0; - int64_t v_idx = vertical_indexes[v++]; - int64_t s_idx = slash_indexes[s++]; - if (causal) { - while (s_idx >= end_m + (kv_seqlen - q_seqlen) && s < NNZ_S) { - s_idx = slash_indexes[s++]; - } - if (s_idx > end_m + (kv_seqlen - q_seqlen)) has_slash = false; - s_idx = max((kv_seqlen - q_seqlen) + end_m - s_idx, BLOCK_SIZE_M); - } else { - while (s_idx >= end_m + kv_seqlen && s < NNZ_S) { - s_idx = slash_indexes[s++]; - } - if (s_idx > end_m + kv_seqlen) has_slash = false; - s_idx = max(kv_seqlen + end_m - s_idx, BLOCK_SIZE_M); - } - - int64_t range_start = s_idx - BLOCK_SIZE_M, range_end = s_idx; - if (!has_slash) { - if (causal) { - range_start = (kv_seqlen - q_seqlen) + end_m; - range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N; - } else { - range_start = kv_seqlen; - range_end = kv_seqlen + BLOCK_SIZE_N; - } - } - - bool slash_finished = false; - while (1) { - if (v_idx < range_end) { - if (v_idx < range_start) { - column_index[tmp_col_cnt++] = v_idx; - } - if (v < NNZ_V) { - v_idx = vertical_indexes[v++]; - } else { - if (causal) - v_idx = end_m + BLOCK_SIZE_N + (kv_seqlen - q_seqlen); - else - v_idx = end_m + BLOCK_SIZE_N + kv_seqlen; - } - } else { - if ((s < NNZ_S && causal) || (s < NNZ_S && !causal && slash_indexes[s] >= start_m)) { - if (causal) - s_idx = max((kv_seqlen - q_seqlen) + end_m - slash_indexes[s++], BLOCK_SIZE_M); - else - s_idx = max(kv_seqlen + end_m - slash_indexes[s++], BLOCK_SIZE_M); - } else { - if (v == NNZ_V || (v_idx > range_start && causal)) { - // add the last vertical if no more slash - if (v == NNZ_V && !causal && v_idx < kv_seqlen) { - column_index[tmp_col_cnt++] = v_idx; - } - tmp_blk_cnt = save_blocks(block_offset, range_start, range_end, BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen); - break; - } else { - if (causal) { - range_start = (kv_seqlen - q_seqlen) + end_m; - range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N; - } else { - // if slash_finished but there are vertical left, save current - // blocks - tmp_blk_cnt = save_blocks(block_offset, range_start, range_end, BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen); - range_start = kv_seqlen; - range_end = kv_seqlen + BLOCK_SIZE_N; - } - slash_finished = true; - } - } - if (!slash_finished) { - if (s_idx > range_end + BLOCK_SIZE_M) { - tmp_blk_cnt = save_blocks(block_offset, range_start, range_end, BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen); - range_start = s_idx - BLOCK_SIZE_M; - range_end = s_idx; - } else if (s_idx > range_end) { - range_end += BLOCK_SIZE_M; - } - } - } - } - - block_count[0] = tmp_blk_cnt; - column_count[0] = tmp_col_cnt; -} - -// Launch the mergehead kernel with 64 threads per block. -void convert_vertical_slash_indexes_64x64_mergehead( - const int* q_seqlens, // [BATCH, ] - const int* kv_seqlens, // [BATCH, ] - const int* vertical_indexes, // [BATCH, N_HEADS, NNZ_V] - const int* slash_indexes, // [BATCH, N_HEADS, NNZ_S] - int* per_head_vertical_topkv, - int* per_head_slash_topkv, - int* block_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)] - int* block_offset, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S] - int* column_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)] - int* column_index, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V] - int64_t BATCH_SIZE, - int64_t N_HEADS, - int64_t N_ROWS, - int64_t BLOCK_SIZE_M, - int64_t BLOCK_SIZE_N, - int64_t NNZ_V, - int64_t NNZ_S, - bool causal) { - const int N_THREADS = 64; - const dim3 dimBlock(N_THREADS); - const dim3 dimGrid(N_HEADS, BATCH_SIZE, (N_ROWS + N_THREADS - 1) / N_THREADS); - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - convert_vertical_slash_indexes_kernel_mergehead<<>>( - q_seqlens, - kv_seqlens, - vertical_indexes, - slash_indexes, - per_head_vertical_topkv, - per_head_slash_topkv, - block_count, - block_offset, - column_count, - column_index, - N_HEADS, - N_ROWS, - BLOCK_SIZE_M, - BLOCK_SIZE_N, - NNZ_V, - NNZ_S, - causal); -} - -// Host wrapper for mergehead kernel. -void convert_vertical_slash_indexes_mergehead( - torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS] - torch::Tensor& block_offset, // [BATCH, N_HEADS, NUM_ROWS, NNZ_S] - torch::Tensor& column_count, // [BATCH, N_HEADS, NUM_ROWS] - torch::Tensor& column_index, // [BATCH, N_HEADS, NUM_ROWS, NNZ_V] - torch::Tensor q_seqlens, // [BATCH, ] - torch::Tensor kv_seqlens, // [BATCH, ] - torch::Tensor vertical_indexes, // [BATCH, N_HEADS, NNZ_V] - torch::Tensor slash_indexes, // [BATCH, N_HEADS, NNZ_S] - torch::Tensor vertical_indices_count, // [N_HEADS, ] - torch::Tensor slash_indices_count, - int64_t context_size, - int64_t block_size_M, - int64_t block_size_N, - bool causal) { - cudaSetDevice(q_seqlens.get_device()); - - int batch_size = slash_indexes.size(0); - int num_heads = slash_indexes.size(1); - int nnz_slash = slash_indexes.size(2); - int nnz_vertical = vertical_indexes.size(2); - int num_rows = (context_size + block_size_M - 1) / block_size_M; - - convert_vertical_slash_indexes_64x64_mergehead( - q_seqlens.data_ptr(), - kv_seqlens.data_ptr(), - vertical_indexes.data_ptr(), - slash_indexes.data_ptr(), - vertical_indices_count.data_ptr(), - slash_indices_count.data_ptr(), - block_count.data_ptr(), - block_offset.data_ptr(), - column_count.data_ptr(), - column_index.data_ptr(), - batch_size, - num_heads, - num_rows, - block_size_M, - block_size_N, - nnz_vertical, - nnz_slash, - causal); -} diff --git a/python/sglang/kernels/aot/csrc/common_extension.cc b/python/sglang/kernels/aot/csrc/common_extension.cc index 211bbdc05..4d0320202 100644 --- a/python/sglang/kernels/aot/csrc/common_extension.cc +++ b/python/sglang/kernels/aot/csrc/common_extension.cc @@ -43,11 +43,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { */ 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( - "cutlass_mla_decode(Tensor! out, Tensor q_nope, Tensor q_pe, Tensor kv_c_and_k_pe_cache, Tensor seq_lens, Tensor " - "page_table, Tensor! workspace, float sm_scale, int num_kv_splits) -> ()"); - m.impl("cutlass_mla_decode", torch::kCUDA, &cutlass_mla_decode); - m.def("cutlass_mla_get_workspace_size", &cutlass_mla_get_workspace_size); /* * From csrc/infllm_v2 @@ -110,9 +105,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { /* * From csrc/gemm */ - m.def("awq_dequantize(Tensor qweight, Tensor scales, Tensor qzeros) -> Tensor"); - m.impl("awq_dequantize", torch::kCUDA, &awq_dequantize); - m.def( "int8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? " "bias) -> Tensor"); @@ -138,17 +130,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { m.def("sgl_per_token_quant_fp8(Tensor input, Tensor! output_q, Tensor! output_s) -> ()"); m.impl("sgl_per_token_quant_fp8", torch::kCUDA, &sgl_per_token_quant_fp8); - /* - * From csrc/gemm/gptq - */ - m.def( - "gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, Tensor b_gptq_scales, Tensor b_g_idx, bool " - "use_shuffle, int bit) -> Tensor"); - m.impl("gptq_gemm", torch::kCUDA, &gptq_gemm); - - m.def("gptq_shuffle(Tensor! q_weight, Tensor q_perm, int bit) -> ()"); - m.impl("gptq_shuffle", torch::kCUDA, &gptq_shuffle); - /* * From csrc/moe */ @@ -353,50 +334,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); - /* - * From Sparse Flash Attention - */ - m.def( - "fwd_sparse(Tensor! q, Tensor k, Tensor v, " - "Tensor block_count, Tensor block_offset, Tensor column_count, Tensor column_index, " - "Tensor!? out, Tensor? alibi_slopes, " - "float p_dropout, float softmax_scale, bool is_causal, " - "float softcap, bool return_softmax, Generator? gen)" - "-> Tensor[]"); - m.impl("fwd_sparse", torch::kCUDA, &flash::mha_fwd_sparse); - - m.def( - "varlen_fwd_sparse(Tensor! q, Tensor k, Tensor v, " - "Tensor block_count, Tensor block_offset, Tensor column_count, Tensor column_index, " - "Tensor!? out, Tensor cu_seqlens_q, " - "Tensor cu_seqlens_k, Tensor? seqused_k, Tensor? alibi_slopes, " - "int max_seqlen_q, int max_seqlen_k, float p_dropout, float softmax_scale, bool zero_tensors, " - "bool is_causal, float softcap, bool return_softmax, " - "Generator? gen) -> Tensor[]"); - m.impl("varlen_fwd_sparse", torch::kCUDA, &flash::mha_varlen_fwd_sparse); - - // Sparse Attention utils - m.def( - "convert_vertical_slash_indexes(" - " Tensor! block_count, Tensor! block_offset, " - " Tensor! column_count, Tensor! column_index, " - " Tensor q_seqlens, Tensor q_seqlens, " - " Tensor vertical_indexes, Tensor slash_indexes, " - " int context_size, int block_size_M, int block_size_N, " - " bool causal) -> ()"); - m.impl("convert_vertical_slash_indexes", torch::kCUDA, &convert_vertical_slash_indexes); - - m.def( - "convert_vertical_slash_indexes_mergehead(" - " Tensor! block_count, Tensor! block_offset, " - " Tensor! column_count, Tensor! column_index, " - " Tensor q_seqlens, Tensor q_seqlens, " - " Tensor vertical_indexes, Tensor slash_indexes, " - " Tensor vertical_indices_count, Tensor slash_indices_count, " - " int context_size, int block_size_M, int block_size_N, " - " bool causal) -> ()"); - m.impl("convert_vertical_slash_indexes_mergehead", torch::kCUDA, &convert_vertical_slash_indexes_mergehead); - /* * From csrc/grammar */ diff --git a/python/sglang/kernels/aot/csrc/common_extension_musa.cc b/python/sglang/kernels/aot/csrc/common_extension_musa.cc index 1c8d71497..d734974a7 100644 --- a/python/sglang/kernels/aot/csrc/common_extension_musa.cc +++ b/python/sglang/kernels/aot/csrc/common_extension_musa.cc @@ -81,9 +81,6 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) { /* * From csrc/gemm */ - m.def("awq_dequantize(Tensor qweight, Tensor scales, Tensor qzeros) -> Tensor"); - m.impl("awq_dequantize", torch::kMUSA, &awq_dequantize); - m.def( "sgl_per_token_group_quant_8bit(Tensor input, Tensor output_q, Tensor output_s, int group_size," " float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()"); diff --git a/python/sglang/kernels/aot/csrc/gemm/awq_kernel.cu b/python/sglang/kernels/aot/csrc/gemm/awq_kernel.cu deleted file mode 100644 index eec933689..000000000 --- a/python/sglang/kernels/aot/csrc/gemm/awq_kernel.cu +++ /dev/null @@ -1,221 +0,0 @@ -// Adapted from -// https://github.com/vllm-project/vllm/blob/eb59b5a6cba6727d3727c0372258db9002f687c1/csrc/quantization/awq/gemm_kernels.cu#L350 -#include -#include -#include -#include -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 -#include -#endif - -template -__device__ inline int lop3(int a, int b, int c) { - int res; - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(res) : "r"(a), "r"(b), "r"(c), "n"(lut)); - return res; -} - -__device__ uint4 dequantize_s4_to_fp16x2(uint32_t const& source) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 750 - uint4 result; - - uint32_t* h = reinterpret_cast(&result); - uint32_t const i4s = reinterpret_cast(source); - - // First, we extract the i4s and construct an intermediate fp16 number. - static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; - static constexpr uint32_t BOTTOM_MASK = 0x000f000f; - static constexpr uint32_t TOP_MASK = 0x00f000f0; - static constexpr uint32_t I4s_TO_F16s_MAGIC_NUM = 0x64006400; - - // Note that the entire sequence only requires 1 shift instruction. This is - // thanks to the register packing format and the fact that we force our - // integers to be unsigned, and account for this in the fp16 subtractions. In - // addition, I exploit the fact that sub and fma have the same throughput in - // order to convert elt_23 and elt_67 to fp16 without having to shift them to - // the bottom bits before hand. - - // Shift right by 8 to now consider elt_45 and elt_67. Issue first to hide RAW - // dependency if we issue immediately before required. - const uint32_t top_i4s = i4s >> 8; - // Extract elt_01 - (i4s & 0x000f000f) | 0x64006400 - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[0]) - : "r"(i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); - // Extract elt_23 (i4s & 0x00f000f0) | 0x64006400 - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[1]) - : "r"(i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); - // Extract elt_45 (top_i4s & 0x000f000f) | 0x64006400 - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[2]) - : "r"(top_i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); - // Extract elt_67 (top_i4s & 0x00f000f0) | 0x64006400 - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[3]) - : "r"(top_i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); - - // This is the half2 {1024, 1024} represented as an integer. - static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64006400; - // This is the half2 {1 / 16, 1 / 16} represented as an integer. - static constexpr uint32_t ONE_SIXTEENTH = 0x2c002c00; - // This is the half2 {-64, -64} represented as an integer. - static constexpr uint32_t NEG_64 = 0xd400d400; - - // Finally, we construct the output numbers. - // Convert elt_01 - asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[0]) : "r"(h[0]), "r"(FP16_TOP_MAGIC_NUM)); - // Convert elt_23 - asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(h[1]) : "r"(h[1]), "r"(ONE_SIXTEENTH), "r"(NEG_64)); - // Convert elt_45 - asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[2]) : "r"(h[2]), "r"(FP16_TOP_MAGIC_NUM)); - // Convert elt_67 - asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(h[3]) : "r"(h[3]), "r"(ONE_SIXTEENTH), "r"(NEG_64)); - - return result; -#else - assert(false); - return {}; -#endif -} - -__device__ uint4 dequantize_s4_to_bf16x2(uint32_t const& source) { -#if CUDA_VERSION >= 12000 -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - uint4 result; - uint32_t* h = reinterpret_cast(&result); - uint32_t const i4s = source; - - // Define masks and constants - static constexpr uint32_t MASK = 0x000f000f; - static constexpr uint32_t EX = 0x43004300; - static constexpr uint32_t MUL = 0x3F803F80; - static constexpr uint32_t ADD = 0xC300C300; - - int lo0 = lop3<(0xf0 & 0xcc) | 0xaa>(i4s, MASK, EX); - int hi0 = lop3<(0xf0 & 0xcc) | 0xaa>(i4s >> 4, MASK, EX); - int lo1 = lop3<(0xf0 & 0xcc) | 0xaa>(i4s >> 8, MASK, EX); - int hi1 = lop3<(0xf0 & 0xcc) | 0xaa>(i4s >> 12, MASK, EX); - - nv_bfloat162* res = reinterpret_cast(h); - res[0] = __hfma2( - *reinterpret_cast(&lo0), - *reinterpret_cast(&MUL), - *reinterpret_cast(&ADD)); - res[1] = __hfma2( - *reinterpret_cast(&hi0), - *reinterpret_cast(&MUL), - *reinterpret_cast(&ADD)); - res[2] = __hfma2( - *reinterpret_cast(&lo1), - *reinterpret_cast(&MUL), - *reinterpret_cast(&ADD)); - res[3] = __hfma2( - *reinterpret_cast(&hi1), - *reinterpret_cast(&MUL), - *reinterpret_cast(&ADD)); - - return result; -#else - assert(false); - return {}; -#endif -#endif -} - -template -__global__ void __launch_bounds__(256) dequantize_weights( - int* __restrict__ qweight, - OutputT* __restrict__ scales, - int* __restrict__ qzeros, - OutputT* __restrict__ output, - int group_size, - int qweight_cols, - int qweight_rows) { -#if CUDA_VERSION >= 12000 - int col = blockIdx.x * blockDim.x + threadIdx.x; - int row = blockIdx.y * blockDim.y + threadIdx.y; - if (col >= qweight_cols || row >= qweight_rows) return; - - int group_idx = row / group_size; - int scale_offset = 8 * col + group_idx * qweight_cols * 8; - uint4 loaded_scale = *(uint4*)(scales + scale_offset); - - // Handle different data types - if constexpr (std::is_same::value) { - // FP16 path - uint4 zeros = dequantize_s4_to_fp16x2(qzeros[col + group_idx * qweight_cols]); - uint4 weight_fp16 = dequantize_s4_to_fp16x2(qweight[col + row * qweight_cols]); - - // Use PTX assembly for FP16 operations - asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.x) : "r"(weight_fp16.x), "r"(zeros.x)); - asm volatile("mul.rn.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.x) : "r"(weight_fp16.x), "r"(loaded_scale.x)); - asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.y) : "r"(weight_fp16.y), "r"(zeros.y)); - asm volatile("mul.rn.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.y) : "r"(weight_fp16.y), "r"(loaded_scale.y)); - asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.z) : "r"(weight_fp16.z), "r"(zeros.z)); - asm volatile("mul.rn.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.z) : "r"(weight_fp16.z), "r"(loaded_scale.z)); - asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.w) : "r"(weight_fp16.w), "r"(zeros.w)); - asm volatile("mul.rn.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.w) : "r"(weight_fp16.w), "r"(loaded_scale.w)); - - OutputT* output_ptr = output + 8 * col + 8 * row * qweight_cols; - *(uint4*)output_ptr = weight_fp16; - } else if constexpr (std::is_same::value) { - uint4 weight_raw = dequantize_s4_to_bf16x2(qweight[col + row * qweight_cols]); - uint4 zero_raw = dequantize_s4_to_bf16x2(qzeros[col + group_idx * qweight_cols]); - uint4 scale_raw = *reinterpret_cast(scales + scale_offset); - - // Vectorized processing (each uint4 contains 4 nv_bfloat162) - nv_bfloat162* weight_vec = reinterpret_cast(&weight_raw); - nv_bfloat162* zero_vec = reinterpret_cast(&zero_raw); - nv_bfloat162* scale_vec = reinterpret_cast(&scale_raw); - -// Single instruction dual-channel operation -#pragma unroll - for (int i = 0; i < 4; ++i) { // uint4 = 4 * nv_bfloat162 - weight_vec[i] = __hmul2(__hsub2(weight_vec[i], zero_vec[i]), scale_vec[i]); - } - - // Directly store to OutputT array (guaranteed contiguous memory) - OutputT* output_ptr = output + 8 * col + row * qweight_cols * 8; - static_assert(sizeof(uint4) == 8 * sizeof(OutputT), "Memory layout mismatch"); - *reinterpret_cast(output_ptr) = weight_raw; - } -#endif -} - -torch::Tensor awq_dequantize(torch::Tensor qweight, torch::Tensor scales, torch::Tensor qzeros) { - int qweight_rows = qweight.size(0); - int qweight_cols = qweight.size(1); - int group_size = qweight_rows / scales.size(0); - - int x_num_threads = 16; - int y_num_threads = 16; - int x_blocks = (qweight_cols + x_num_threads - 1) / x_num_threads; - int y_blocks = (qweight_rows + y_num_threads - 1) / y_num_threads; - - const at::cuda::OptionalCUDAGuard device_guard(device_of(qweight)); - - auto output_tensor_options = torch::TensorOptions().dtype(scales.dtype()).device(scales.device()); - at::Tensor output = torch::empty({qweight_rows, qweight_cols * 8}, output_tensor_options); - - auto _qweight = reinterpret_cast(qweight.data_ptr()); - auto _zeros = reinterpret_cast(qzeros.data_ptr()); - - dim3 num_blocks(x_blocks, y_blocks); - dim3 threads_per_block(x_num_threads, y_num_threads); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - - if (scales.scalar_type() == at::ScalarType::Half) { - auto _scales = reinterpret_cast(scales.data_ptr()); - auto _output = reinterpret_cast(output.data_ptr()); - dequantize_weights<<>>( - _qweight, _scales, _zeros, _output, group_size, qweight_cols, qweight_rows); - } else { - auto _scales = reinterpret_cast<__nv_bfloat16*>(scales.data_ptr()); - auto _output = reinterpret_cast<__nv_bfloat16*>(output.data_ptr()); - dequantize_weights<__nv_bfloat16><<>>( - _qweight, _scales, _zeros, _output, group_size, qweight_cols, qweight_rows); - } - - return output; -} diff --git a/python/sglang/kernels/aot/csrc/gemm/gptq/compat.cuh b/python/sglang/kernels/aot/csrc/gemm/gptq/compat.cuh deleted file mode 100644 index 506eeb769..000000000 --- a/python/sglang/kernels/aot/csrc/gemm/gptq/compat.cuh +++ /dev/null @@ -1,62 +0,0 @@ -/* -Copied from https://github.com/turboderp/exllamav2 -*/ - -#ifndef _compat_cuh -#define _compat_cuh - -namespace sglang { -namespace gptq { -// atomicAdd for half types, to support CC < 7.x - -__device__ __forceinline__ void atomicAdd_half(half* address, half val) { - unsigned int* address_as_ui = (unsigned int*)((char*)address - ((size_t)address & 2)); - unsigned int old = *address_as_ui; - unsigned int assumed; - - do { - assumed = old; - __half_raw hsum; - hsum.x = (size_t)address & 2 ? (old >> 16) : (old & 0xffff); - half tmpres = __hadd(hsum, val); - hsum = __half_raw(tmpres); - old = (size_t)address & 2 ? (old & 0xffff) | (hsum.x << 16) : (old & 0xffff0000) | hsum.x; - old = atomicCAS(address_as_ui, assumed, old); - } while (assumed != old); -} - -// atomicAdd for half2 types - -__device__ __forceinline__ void atomicAdd_half2(half2* address, half2 val) { - unsigned int* address_as_ui = (unsigned int*)address; - unsigned int old = *address_as_ui; - unsigned int assumed; - do { - assumed = old; - half2 old_val = *((half2*)&old); - half2 new_val = __hadd2(old_val, val); - old = atomicCAS(address_as_ui, assumed, *((unsigned int*)&new_val)); - } while (assumed != old); -} - -// - -#if defined(__CUDA_ARCH__) || defined(USE_ROCM) -#if __CUDA_ARCH__ < 700 || defined(USE_ROCM) - -__device__ __forceinline__ void atomicAdd(half* address, half val) { - atomicAdd_half(address, val); -} - -#if __CUDA_ARCH__ < 600 || defined(USE_ROCM) -__device__ __forceinline__ void atomicAdd(half2* address, half2 val) { - atomicAdd_half2(address, val); -} -#endif - -#endif -#endif - -} // namespace gptq -} // namespace sglang -#endif diff --git a/python/sglang/kernels/aot/csrc/gemm/gptq/gptq_kernel.cu b/python/sglang/kernels/aot/csrc/gemm/gptq/gptq_kernel.cu deleted file mode 100644 index 4dd5d8a24..000000000 --- a/python/sglang/kernels/aot/csrc/gemm/gptq/gptq_kernel.cu +++ /dev/null @@ -1,1950 +0,0 @@ -/* -Adapted from https://github.com/turboderp/exllamav2 and -https://github.com/qwopqwop200/GPTQ-for-LLaMa -*/ - -#include -#include -#include -#include -#include - -#include -#include - -#include "compat.cuh" -#include "matrix_view.cuh" -#include "qdq_2.cuh" -#include "qdq_3.cuh" -#include "qdq_4.cuh" -#include "qdq_8.cuh" - -namespace sglang { -namespace gptq { - -#define BLOCK_KN_SIZE 128 -#define BLOCK_M_SIZE_MAX 8 -#define MAX_GROUPS_IN_BLOCK (BLOCK_KN_SIZE / 32) -#define MAX_Q_GEMM_ROWS 50 -#define MAX_Q_GEMM_ROWS_8BIT 24 -#define MAX_ALT_GEMM_ROWS 8 -#define THREADS_X 32 -#define THREADS_Y 32 -#define DIVIDE(x, size) (((x) + (size) - 1) / (size)) - -#if defined(USE_ROCM) -#include -__host__ __forceinline__ hipblasStatus_t __compat_hipblasHgemm( - hipblasHandle_t handle, - hipblasOperation_t transA, - hipblasOperation_t transB, - int m, - int n, - int k, - const half* alpha, - const half* AP, - int lda, - const half* BP, - int ldb, - const half* beta, - half* CP, - int ldc) { - return hipblasHgemm( - handle, - transA, - transB, - m, - n, - k, - reinterpret_cast(alpha), - reinterpret_cast(AP), - lda, - reinterpret_cast(BP), - ldb, - reinterpret_cast(beta), - reinterpret_cast(CP), - ldc); -} -#define hipblasHgemm __compat_hipblasHgemm - -// Previous version of PyTorch were converting to rocBLAS instead of hipBLAS. -#define rocblas_operation_none HIPBLAS_OP_N -#define rocblas_hgemm __compat_hipblasHgemm -#endif - -__forceinline__ __device__ half2 dot22_8(half2 (&dq)[4], const half* a_ptr, const half2 g_result) { - half2 result = {}; - const half2* a2_ptr = (const half2*)a_ptr; -#pragma unroll - for (int i = 0; i < 4; i++) - result = __hfma2(dq[i], *a2_ptr++, result); - return __hadd2(result, g_result); -} - -__forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) { - half2 result = {}; - const half2* a2_ptr = (const half2*)a_ptr; -#pragma unroll - for (int i = 0; i < 4; i++) - result = __hfma2(dq[i], *a2_ptr++, result); - return __half2float(__low2half(result)) + __half2float(__high2half(result)); -} - -__forceinline__ __device__ half2 dot22_8(half2 (&dq)[4], const half* a_ptr, const half2 g_result, const half qs_h) { - half2 result = {}; - const half2* a2_ptr = (const half2*)a_ptr; -#pragma unroll - for (int i = 0; i < 4; i++) - result = __hfma2(dq[i], *a2_ptr++, result); - return __hfma2(result, __halves2half2(qs_h, qs_h), g_result); -} - -__forceinline__ __device__ half2 dot22_16(half2 (&dq)[8], const half* a_ptr, const half2 g_result, const half qs_h) { - half2 result = {}; - const half2* a2_ptr = (const half2*)a_ptr; -#pragma unroll - for (int i = 0; i < 8; i++) - result = __hfma2(dq[i], *a2_ptr++, result); - return __hfma2(result, __halves2half2(qs_h, qs_h), g_result); -} - -__forceinline__ __device__ half2 dot22_32(half2 (&dq)[16], const half* a_ptr, const half2 g_result, const half qs_h) { - half2 result = {}; - const half2* a2_ptr = (const half2*)a_ptr; -#pragma unroll - for (int i = 0; i < 16; i += 1) - result = __hfma2(dq[i], *a2_ptr++, result); - return __hfma2(result, __halves2half2(qs_h, qs_h), g_result); -} - -__forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr, const float g_result, const float qs_f) { - half2 result = {}; - const half2* a2_ptr = (const half2*)a_ptr; -#pragma unroll - for (int i = 0; i < 4; i++) - result = __hfma2(dq[i], *a2_ptr++, result); - float result_f = __half2float(__low2half(result)) + __half2float(__high2half(result)); - return fma(result_f, qs_f, g_result); -} - -__forceinline__ __device__ float dot22_16_f(half2 (&dq)[8], const half* a_ptr, const float g_result, const float qs_f) { - half2 result = {}; - const half2* a2_ptr = (const half2*)a_ptr; -#pragma unroll - for (int i = 0; i < 8; i++) - result = __hfma2(dq[i], *a2_ptr++, result); - float result_f = __half2float(__low2half(result)) + __half2float(__high2half(result)); - return fma(result_f, qs_f, g_result); -} - -__forceinline__ __device__ float -dot22_32_f(half2 (&dq)[16], const half* a_ptr, const float g_result, const float qs_f) { - half2 result = {}; - const half2* a2_ptr = (const half2*)a_ptr; -#pragma unroll - for (int i = 0; i < 16; i += 1) - result = __hfma2(dq[i], *a2_ptr++, result); - float result_f = __half2float(__low2half(result)) + __half2float(__high2half(result)); - return fma(result_f, qs_f, g_result); -} - -__forceinline__ __device__ half dot22_8_h(half2 (&dq)[4], const half* a_ptr, const half g_result, const half qs_h) { - // Use FP32 accumulator to avoid potential overflow since unscaled weights are - // in the range -128..127 - - float result = {}; -#pragma unroll - for (int i = 0; i < 4; i++) { - half2 w01 = dq[i]; - float w0 = __low2float(w01); - float w1 = __high2float(w01); - float x0 = __half2float(*a_ptr++); - float x1 = __half2float(*a_ptr++); - result = fma(w0, x0, result); - result = fma(w1, x1, result); - } - float qs = __half2float(qs_h); - result *= qs; - half result_h = __float2half_rn(result); - return __hadd(result_h, g_result); -} - -__forceinline__ __device__ half dot22_16_h(half2 (&dq)[8], const half* a_ptr, const half g_result, const half qs_h) { - half2 result = {}; - const half2* a2_ptr = (const half2*)a_ptr; -#pragma unroll - for (int i = 0; i < 8; i++) - result = __hfma2(dq[i], *a2_ptr++, result); - half result_h = __hadd(__low2half(result), __high2half(result)); - return __hfma(result_h, qs_h, g_result); -} - -__forceinline__ __device__ half dot22_32_h(half2 (&dq)[16], const half* a_ptr, const half g_result, const half qs_h) { - half2 result = {}; - const half2* a2_ptr = (const half2*)a_ptr; -#pragma unroll - for (int i = 0; i < 16; i += 1) - result = __hfma2(dq[i], *a2_ptr++, result); - half result_h = __hadd(__low2half(result), __high2half(result)); - return __hfma(result_h, qs_h, g_result); -} - -typedef void (*fp_gemm_half_q_half_gptq_kernel)( - const half*, - const uint32_t*, - const uint32_t*, - const half*, - half*, - const int, - const int, - const int, - const int, - const int*); - -template -__global__ void gemm_half_q_half_gptq_4bit_kernel( - const half* __restrict__ a, - const uint32_t* __restrict__ b_q_weight, - const uint32_t* __restrict__ b_gptq_qzeros, - const half* __restrict__ b_gptq_scales, - half* __restrict__ c, - const int size_m, - const int size_n, - const int size_k, - const int groups, - const int* __restrict__ b_q_perm) { - MatrixView_half a_(a, size_m, size_k); - MatrixView_half_rw c_(c, size_m, size_n); - MatrixView_q4_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n); - MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n); - - auto t = threadIdx.x; - - // Block - auto offset_n = blockIdx.x * BLOCK_KN_SIZE * 4; - auto offset_m = blockIdx.y * m_count; - auto offset_k = blockIdx.z * BLOCK_KN_SIZE; - - int end_k = min(offset_k + BLOCK_KN_SIZE, size_k); - - int n = offset_n + t * 4; - - // Preload block_a - __shared__ half block_a[m_count][BLOCK_KN_SIZE]; - - if (offset_k + t < end_k) { - for (int m = 0; m < m_count; ++m) { - const half* a_ptr = a_.item_ptr(offset_m + m, 0); - half* block_a_ptr = block_a[m]; - - half a0; - if (b_q_perm) - a0 = a_ptr[b_q_perm[offset_k + t]]; - else - a0 = a_ptr[offset_k + t]; - block_a_ptr[t] = a0; - } - } - - // Zero output - if (n >= size_n) return; - - if (blockIdx.z == 0) { - for (int m = 0; m < m_count; m++) - *((uint64_t*)c_.item_ptr(offset_m + m, n)) = 0; - } - - __syncthreads(); - - // Find initial group - int groupsize = size_k / groups; - int group = offset_k / groupsize; - int nextgroup = offset_k + groupsize; - - // a, b offset - int qk = offset_k / (32 / 4); - - const uint32_t* b_ptr = b_q_weight + qk * size_n + n; - const half* a_ptr = &block_a[0][0]; - int a_stride = BLOCK_KN_SIZE; - - // Initial group - int zeros[4]; - float scales[4]; - half2 z1z16[4][2]; - half2 y1y16[4][2]; - b_gptq_qzeros_.item4(zeros, group, n); - b_gptq_scales_.item4_f(scales, group, n); - dequant_4bit_8_prep_zero(zeros[0] + 1, z1z16[0], y1y16[0]); - dequant_4bit_8_prep_zero(zeros[1] + 1, z1z16[1], y1y16[1]); - dequant_4bit_8_prep_zero(zeros[2] + 1, z1z16[2], y1y16[2]); - dequant_4bit_8_prep_zero(zeros[3] + 1, z1z16[3], y1y16[3]); - - // Column result - float block_c[m_count][4] = {}; - - // Dequantize and multiply - int k = offset_k; - while (k < end_k) { - if (k == nextgroup) { - group++; - nextgroup += groupsize; - b_gptq_qzeros_.item4(zeros, group, n); - b_gptq_scales_.item4_f(scales, group, n); - dequant_4bit_8_prep_zero(zeros[0] + 1, z1z16[0], y1y16[0]); - dequant_4bit_8_prep_zero(zeros[1] + 1, z1z16[1], y1y16[1]); - dequant_4bit_8_prep_zero(zeros[2] + 1, z1z16[2], y1y16[2]); - dequant_4bit_8_prep_zero(zeros[3] + 1, z1z16[3], y1y16[3]); - } - -#pragma unroll - for (int j = 0; j < 4; j++) { - const int4* b_ptr4 = (int4*)b_ptr; - int4 load_int4 = *b_ptr4; - - half2 dq[4][4]; - dequant_4bit_8_gptq(load_int4.x, dq[0], z1z16[0], y1y16[0], size_n, false); - dequant_4bit_8_gptq(load_int4.y, dq[1], z1z16[1], y1y16[1], size_n, false); - dequant_4bit_8_gptq(load_int4.z, dq[2], z1z16[2], y1y16[2], size_n, false); - dequant_4bit_8_gptq(load_int4.w, dq[3], z1z16[3], y1y16[3], size_n, false); - -#pragma unroll - for (int m = 0; m < m_count; m++) { - block_c[m][0] = fma(dot22_8_f(dq[0], a_ptr + m * a_stride), scales[0], block_c[m][0]); - block_c[m][1] = fma(dot22_8_f(dq[1], a_ptr + m * a_stride), scales[1], block_c[m][1]); - block_c[m][2] = fma(dot22_8_f(dq[2], a_ptr + m * a_stride), scales[2], block_c[m][2]); - block_c[m][3] = fma(dot22_8_f(dq[3], a_ptr + m * a_stride), scales[3], block_c[m][3]); - } - - b_ptr += size_n; - a_ptr += 8; - } - - k += 32; - } - - for (int m = 0; m < m_count; m++) { - half2* out = (half2*)c_.item_ptr(offset_m + m, n); - half2 result01 = __halves2half2(__float2half_rn(block_c[m][0]), __float2half_rn(block_c[m][1])); - half2 result23 = __halves2half2(__float2half_rn(block_c[m][2]), __float2half_rn(block_c[m][3])); - atomicAdd(out, result01); - atomicAdd(out + 1, result23); - } -} - -template -__global__ void gemm_half_q_half_gptq_2bit_kernel( - const half* __restrict__ a, - const uint32_t* __restrict__ b_q_weight, - const uint32_t* __restrict__ b_gptq_qzeros, - const half* __restrict__ b_gptq_scales, - half* __restrict__ c, - const int size_m, - const int size_n, - const int size_k, - const int groups, - const int* __restrict__ b_q_perm) { - MatrixView_half a_(a, size_m, size_k); - MatrixView_half_rw c_(c, size_m, size_n); - MatrixView_q2_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n); - MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n); - - auto t = threadIdx.x; - - // Block - auto offset_n = blockIdx.x * BLOCK_KN_SIZE * 4; - auto offset_m = blockIdx.y * m_count; - auto offset_k = blockIdx.z * BLOCK_KN_SIZE; - - int end_k = min(offset_k + BLOCK_KN_SIZE, size_k); - - int n = offset_n + t * 4; - - // Preload block_a - __shared__ half block_a[m_count][BLOCK_KN_SIZE]; - - if (offset_k + t < end_k) { - for (int m = 0; m < m_count; ++m) { - const half* a_ptr = a_.item_ptr(offset_m + m, 0); - half* block_a_ptr = block_a[m]; - - half a0; - if (b_q_perm) - a0 = a_ptr[b_q_perm[offset_k + t]]; - else - a0 = a_ptr[offset_k + t]; - block_a_ptr[t] = a0; - } - } - - // Zero output - if (n >= size_n) return; - - if (blockIdx.z == 0) { - for (int m = 0; m < m_count; m++) - *((uint64_t*)c_.item_ptr(offset_m + m, n)) = 0; - } - - __syncthreads(); - - // Find initial group - int groupsize = size_k / groups; - int group = offset_k / groupsize; - int nextgroup = offset_k + groupsize; - - // a, b offset - int qk = offset_k / (32 / 2); - - const uint32_t* b_ptr = b_q_weight + qk * size_n + n; - const half* a_ptr = &block_a[0][0]; - int a_stride = BLOCK_KN_SIZE; - - // Initial group - int zeros[4]; - half scales[4]; - b_gptq_qzeros_.item4(zeros, group, n); - b_gptq_scales_.item4(scales, group, n); - // Column result - half block_c[m_count][4] = {}; - - // Dequantize and multiply - int k = offset_k; - while (k < end_k) { - if (k == nextgroup) { - group++; - nextgroup += groupsize; - b_gptq_qzeros_.item4(zeros, group, n); - b_gptq_scales_.item4(scales, group, n); - } - -#pragma unroll - for (int j = 0; j < 1; j++) { - const int4* b_ptr4 = (int4*)b_ptr; - int4 load_int4 = *b_ptr4; - - half2 dq[4][8]; - dequant_2bit_16(load_int4.x, dq[0], size_n, zeros[0] + 1); - dequant_2bit_16(load_int4.y, dq[1], size_n, zeros[1] + 1); - dequant_2bit_16(load_int4.z, dq[2], size_n, zeros[2] + 1); - dequant_2bit_16(load_int4.w, dq[3], size_n, zeros[3] + 1); - -#pragma unroll - for (int m = 0; m < m_count; m++) { - block_c[m][0] = dot22_16_h(dq[0], a_ptr + m * a_stride, block_c[m][0], scales[0]); - block_c[m][1] = dot22_16_h(dq[1], a_ptr + m * a_stride, block_c[m][1], scales[1]); - block_c[m][2] = dot22_16_h(dq[2], a_ptr + m * a_stride, block_c[m][2], scales[2]); - block_c[m][3] = dot22_16_h(dq[3], a_ptr + m * a_stride, block_c[m][3], scales[3]); - } - - b_ptr += size_n; - a_ptr += 16; - } - - k += 16; - } - - for (int m = 0; m < m_count; m++) { - half2* out = (half2*)c_.item_ptr(offset_m + m, n); - half2 result01 = __halves2half2(block_c[m][0], block_c[m][1]); - half2 result23 = __halves2half2(block_c[m][2], block_c[m][3]); - atomicAdd(out, result01); - atomicAdd(out + 1, result23); - } -} - -template -__global__ void gemm_half_q_half_gptq_3bit_kernel( - const half* __restrict__ a, - const uint32_t* __restrict__ b_q_weight, - const uint32_t* __restrict__ b_gptq_qzeros, - const half* __restrict__ b_gptq_scales, - half* __restrict__ c, - const int size_m, - const int size_n, - const int size_k, - const int groups, - const int* __restrict__ b_q_perm) { - MatrixView_half a_(a, size_m, size_k); - MatrixView_half_rw c_(c, size_m, size_n); - MatrixView_q3_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n); - MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n); - - auto t = threadIdx.x; - - // Block - auto offset_n = blockIdx.x * BLOCK_KN_SIZE * 4; - auto offset_m = blockIdx.y * m_count; - auto offset_k = blockIdx.z * BLOCK_KN_SIZE; - - int end_k = min(offset_k + BLOCK_KN_SIZE, size_k); - - int n = offset_n + t * 4; - - // Preload block_a - __shared__ half block_a[m_count][BLOCK_KN_SIZE]; - - if (offset_k + t < end_k) { - for (int m = 0; m < m_count; ++m) { - const half* a_ptr = a_.item_ptr(offset_m + m, 0); - half* block_a_ptr = block_a[m]; - - half a0; - if (b_q_perm) - a0 = a_ptr[b_q_perm[offset_k + t]]; - else - a0 = a_ptr[offset_k + t]; - block_a_ptr[t] = a0; - } - } - - // Zero output - if (n >= size_n) return; - - if (blockIdx.z == 0) { - for (int m = 0; m < m_count; m++) - *((uint64_t*)c_.item_ptr(offset_m + m, n)) = 0; - } - - __syncthreads(); - - // Find initial group - int groupsize = size_k / groups; - int group = offset_k / groupsize; - int nextgroup = offset_k + groupsize; - - // a, b offset - int qk = offset_k / 32 * 3; - - const uint32_t* b_ptr = b_q_weight + qk * size_n + n; - const half* a_ptr = &block_a[0][0]; - int a_stride = BLOCK_KN_SIZE; - - // Initial group - int zeros[4]; - half scales[4]; - b_gptq_qzeros_.item4(zeros, group, n); - b_gptq_scales_.item4(scales, group, n); - // Column result - half block_c[m_count][4] = {}; - - // Dequantize and multiply - int k = offset_k; - while (k < end_k) { - if (k == nextgroup) { - group++; - nextgroup += groupsize; - b_gptq_qzeros_.item4(zeros, group, n); - b_gptq_scales_.item4(scales, group, n); - } - -#pragma unroll - for (int j = 0; j < 1; j++) { - int4 load_int4[3]; - load_int4[0] = *((int4*)b_ptr); - b_ptr += size_n; - load_int4[1] = *((int4*)b_ptr); - b_ptr += size_n; - load_int4[2] = *((int4*)b_ptr); - b_ptr += size_n; - - half2 dq[4][16]; - dequant_3bit_32(load_int4[0].x, load_int4[1].x, load_int4[2].x, dq[0], size_n, zeros[0] + 1); - dequant_3bit_32(load_int4[0].y, load_int4[1].y, load_int4[2].y, dq[1], size_n, zeros[1] + 1); - dequant_3bit_32(load_int4[0].z, load_int4[1].z, load_int4[2].z, dq[2], size_n, zeros[2] + 1); - dequant_3bit_32(load_int4[0].w, load_int4[1].w, load_int4[2].w, dq[3], size_n, zeros[3] + 1); - -#pragma unroll - for (int m = 0; m < m_count; m++) { - block_c[m][0] = dot22_32_h(dq[0], a_ptr + m * a_stride, block_c[m][0], scales[0]); - block_c[m][1] = dot22_32_h(dq[1], a_ptr + m * a_stride, block_c[m][1], scales[1]); - block_c[m][2] = dot22_32_h(dq[2], a_ptr + m * a_stride, block_c[m][2], scales[2]); - block_c[m][3] = dot22_32_h(dq[3], a_ptr + m * a_stride, block_c[m][3], scales[3]); - } - a_ptr += 32; - } - - k += 32; - } - - for (int m = 0; m < m_count; m++) { - half2* out = (half2*)c_.item_ptr(offset_m + m, n); - half2 result01 = __halves2half2(block_c[m][0], block_c[m][1]); - half2 result23 = __halves2half2(block_c[m][2], block_c[m][3]); - atomicAdd(out, result01); - atomicAdd(out + 1, result23); - } -} - -template -__global__ void gemm_half_q_half_gptq_8bit_kernel( - const half* __restrict__ a, - const uint32_t* __restrict__ b_q_weight, - const uint32_t* __restrict__ b_gptq_qzeros, - const half* __restrict__ b_gptq_scales, - half* __restrict__ c, - const int size_m, - const int size_n, - const int size_k, - const int groups, - const int* __restrict__ b_q_perm) { - MatrixView_half a_(a, size_m, size_k); - MatrixView_half_rw c_(c, size_m, size_n); - MatrixView_q8_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n); - MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n); - - auto t = threadIdx.x; - - // Block - auto offset_n = blockIdx.x * BLOCK_KN_SIZE * 4; - auto offset_m = blockIdx.y * m_count; - auto offset_k = blockIdx.z * BLOCK_KN_SIZE; - - int end_k = min(offset_k + BLOCK_KN_SIZE, size_k); - - int n = offset_n + t * 4; - - // Preload block_a - __shared__ half block_a[m_count][BLOCK_KN_SIZE]; - - if (offset_k + t < end_k) { - for (int m = 0; m < m_count; ++m) { - const half* a_ptr = a_.item_ptr(offset_m + m, 0); - half* block_a_ptr = block_a[m]; - - half a0; - if (b_q_perm) - a0 = a_ptr[b_q_perm[offset_k + t]]; - else - a0 = a_ptr[offset_k + t]; - block_a_ptr[t] = a0; - } - } - - // Zero output - if (n >= size_n) return; - - if (blockIdx.z == 0) { - for (int m = 0; m < m_count; m++) - *((uint64_t*)c_.item_ptr(offset_m + m, n)) = 0; - } - - __syncthreads(); - - // Find initial group - int groupsize = size_k / groups; - int group = offset_k / groupsize; - int nextgroup = offset_k + groupsize; - - // a, b offset - int qk = offset_k / (32 / 8); - - const uint32_t* b_ptr = b_q_weight + qk * size_n + n; - const half* a_ptr = &block_a[0][0]; - int a_stride = BLOCK_KN_SIZE; - - // Initial group - int zeros[4]; - half scales[4]; - b_gptq_qzeros_.item4(zeros, group, n); - b_gptq_scales_.item4(scales, group, n); - // Column result - half block_c[m_count][4] = {}; - - // Dequantize and multiply - int k = offset_k; - while (k < end_k) { - if (k == nextgroup) { - group++; - nextgroup += groupsize; - b_gptq_qzeros_.item4(zeros, group, n); - b_gptq_scales_.item4(scales, group, n); - } - -#pragma unroll - for (int j = 0; j < 4; j++) { - int4 load_int4[2]; - load_int4[0] = *((int4*)b_ptr); - b_ptr += size_n; - load_int4[1] = *((int4*)b_ptr); - b_ptr += size_n; - - half2 dq[4][4]; - dequant_8bit_8(load_int4[0].x, load_int4[1].x, dq[0], size_n, zeros[0] + 1); - dequant_8bit_8(load_int4[0].y, load_int4[1].y, dq[1], size_n, zeros[1] + 1); - dequant_8bit_8(load_int4[0].z, load_int4[1].z, dq[2], size_n, zeros[2] + 1); - dequant_8bit_8(load_int4[0].w, load_int4[1].w, dq[3], size_n, zeros[3] + 1); - - for (int m = 0; m < m_count; m++) { - block_c[m][0] = dot22_8_h(dq[0], a_ptr + m * a_stride, block_c[m][0], scales[0]); - block_c[m][1] = dot22_8_h(dq[1], a_ptr + m * a_stride, block_c[m][1], scales[1]); - block_c[m][2] = dot22_8_h(dq[2], a_ptr + m * a_stride, block_c[m][2], scales[2]); - block_c[m][3] = dot22_8_h(dq[3], a_ptr + m * a_stride, block_c[m][3], scales[3]); - } - a_ptr += 8; - } - k += 32; - } - - for (int m = 0; m < m_count; m++) { - half2* out = (half2*)c_.item_ptr(offset_m + m, n); - half2 result01 = __halves2half2(block_c[m][0], block_c[m][1]); - half2 result23 = __halves2half2(block_c[m][2], block_c[m][3]); - atomicAdd(out, result01); - atomicAdd(out + 1, result23); - } -} - -fp_gemm_half_q_half_gptq_kernel pick_gemm_half_q_half_gptq_kernel(bool first_block, const int m_count, const int bit) { -#define SELECT_KERNEL(M_COUNT) \ - if (m_count == M_COUNT) { \ - if (bit == 2) return gemm_half_q_half_gptq_2bit_kernel; \ - if (bit == 3) return gemm_half_q_half_gptq_3bit_kernel; \ - if (bit == 4) return gemm_half_q_half_gptq_4bit_kernel; \ - if (bit == 8) return gemm_half_q_half_gptq_8bit_kernel; \ - } -#if BLOCK_M_SIZE_MAX >= 1 - SELECT_KERNEL(1); -#endif -#if BLOCK_M_SIZE_MAX >= 2 - SELECT_KERNEL(2); -#endif -#if BLOCK_M_SIZE_MAX >= 3 - SELECT_KERNEL(3); -#endif -#if BLOCK_M_SIZE_MAX >= 4 - SELECT_KERNEL(4); -#endif -#if BLOCK_M_SIZE_MAX >= 5 - SELECT_KERNEL(5); -#endif -#if BLOCK_M_SIZE_MAX >= 6 - SELECT_KERNEL(6); -#endif -#if BLOCK_M_SIZE_MAX >= 7 - SELECT_KERNEL(7); -#endif -#if BLOCK_M_SIZE_MAX >= 8 - SELECT_KERNEL(8); -#endif - return NULL; -} - -void gemm_half_q_half_cuda_part( - const half* a, - const uint32_t* b_q_weight, - const uint32_t* b_gptq_qzeros, - const half* b_gptq_scales, - const int* b_q_perm, - half* c, - int size_m, - int size_n, - int size_k, - int m_count, - int groups, - int bit) { - dim3 blockDim, gridDim; - blockDim.x = BLOCK_KN_SIZE; - blockDim.y = 1; - blockDim.z = 1; - gridDim.x = DIVIDE(size_n, BLOCK_KN_SIZE * 4); - gridDim.y = DIVIDE(size_m, m_count); - gridDim.z = DIVIDE(size_k, BLOCK_KN_SIZE); - - fp_gemm_half_q_half_gptq_kernel kernel = pick_gemm_half_q_half_gptq_kernel(true, m_count, bit); - - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - kernel<<>>( - a, b_q_weight, b_gptq_qzeros, b_gptq_scales, c, size_m, size_n, size_k, groups, b_q_perm); -} - -__global__ void reconstruct_exllama_8bit_kernel( - const uint32_t* __restrict__ b_q_weight, - const int* __restrict__ b_q_perm, - const uint32_t* __restrict__ b_gptq_qzeros, - const half* __restrict__ b_gptq_scales, - const int size_k, - const int size_n, - const int groups, - half* __restrict__ b) { - MatrixView_half_rw b_(b, size_k, size_n); - MatrixView_q8_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n); - MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n); - - auto offset_k = BLOCK_KN_SIZE * blockIdx.y; - auto offset_n = BLOCK_KN_SIZE * blockIdx.x * 4; - - int end_k = min(offset_k + BLOCK_KN_SIZE, size_k); - - // Preload remapping table - __shared__ int perm[BLOCK_KN_SIZE]; - auto t = threadIdx.x; - - if (b_q_perm) { - if (offset_k + t < size_k) perm[t] = b_q_perm[offset_k + t]; - } - - // Column - int n = offset_n + t * 4; - if (n >= size_n) return; - - // Find initial group - int groupsize = size_k / groups; - int group = offset_k / groupsize; - int nextgroup = offset_k + groupsize; - - // b offset - int qk = offset_k / (32 / 8); - - const uint32_t* b_ptr = b_q_weight + qk * size_n + n; - - // Initial zeros/scale - int zeros[4]; - half2 scales[4]; - b_gptq_qzeros_.item4(zeros, group, n); - b_gptq_scales_.item4_h2(scales, group, n); - - __syncthreads(); - - int k = offset_k; - int lk = 0; - - while (k < end_k) { - if (k == nextgroup) { - group++; - nextgroup += groupsize; - b_gptq_qzeros_.item4(zeros, group, n); - b_gptq_scales_.item4_h2(scales, group, n); - } - - for (int p = 0; p < 4; p++) { - int4 load_int4[2]; - load_int4[0] = *((int4*)b_ptr); - b_ptr += size_n; - load_int4[1] = *((int4*)b_ptr); - b_ptr += size_n; - - half2 dq[4][4]; - dequant_8bit_8(load_int4[0].x, load_int4[1].x, dq[0], size_n, zeros[0] + 1); - dequant_8bit_8(load_int4[0].y, load_int4[1].y, dq[1], size_n, zeros[1] + 1); - dequant_8bit_8(load_int4[0].z, load_int4[1].z, dq[2], size_n, zeros[2] + 1); - dequant_8bit_8(load_int4[0].w, load_int4[1].w, dq[3], size_n, zeros[3] + 1); - - // half* dqh = (half*)dq; - if (b_q_perm) { - for (int j = 0; j < 4; j++) { - for (int v = 0; v < 4; v++) - dq[v][j] = __hmul2(scales[v], dq[v][j]); - b_.set4( - perm[lk++], n, __low2half(dq[0][j]), __low2half(dq[1][j]), __low2half(dq[2][j]), __low2half(dq[3][j])); - b_.set4( - perm[lk++], - n, - __high2half(dq[0][j]), - __high2half(dq[1][j]), - __high2half(dq[2][j]), - __high2half(dq[3][j])); - } - } else { - for (int j = 0; j < 4; j++) { - for (int v = 0; v < 4; v++) - dq[v][j] = __hmul2(scales[v], dq[v][j]); - b_.set4( - offset_k + lk++, - n, - __low2half(dq[0][j]), - __low2half(dq[1][j]), - __low2half(dq[2][j]), - __low2half(dq[3][j])); - b_.set4( - offset_k + lk++, - n, - __high2half(dq[0][j]), - __high2half(dq[1][j]), - __high2half(dq[2][j]), - __high2half(dq[3][j])); - } - } - } - k += 32; - } -} - -__global__ void reconstruct_exllama_4bit_kernel( - const uint32_t* __restrict__ b_q_weight, - const int* __restrict__ b_q_perm, - const uint32_t* __restrict__ b_gptq_qzeros, - const half* __restrict__ b_gptq_scales, - const int size_k, - const int size_n, - const int groups, - half* __restrict__ b) { - MatrixView_half_rw b_(b, size_k, size_n); - MatrixView_q4_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n); - MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n); - - auto offset_k = BLOCK_KN_SIZE * blockIdx.y; - auto offset_n = BLOCK_KN_SIZE * blockIdx.x * 4; - - int end_k = min(offset_k + BLOCK_KN_SIZE, size_k); - - // Preload remapping table - __shared__ int perm[BLOCK_KN_SIZE]; - auto t = threadIdx.x; - - if (b_q_perm) { - if (offset_k + t < size_k) perm[t] = b_q_perm[offset_k + t]; - } - - // Column - int n = offset_n + t * 4; - if (n >= size_n) return; - - // Find initial group - int groupsize = size_k / groups; - int group = offset_k / groupsize; - int nextgroup = offset_k + groupsize; - - // b offset - int qk = offset_k / (32 / 4); - - const uint32_t* b_ptr = b_q_weight + qk * size_n + n; - - // Initial zeros/scale - int zeros[4]; - half2 scales[4]; - half2 z1z16[4][2]; - half2 y1y16[4][2]; - b_gptq_qzeros_.item4(zeros, group, n); - b_gptq_scales_.item4_h2(scales, group, n); - dequant_4bit_8_prep_zero(zeros[0] + 1, z1z16[0], y1y16[0]); - dequant_4bit_8_prep_zero(zeros[1] + 1, z1z16[1], y1y16[1]); - dequant_4bit_8_prep_zero(zeros[2] + 1, z1z16[2], y1y16[2]); - dequant_4bit_8_prep_zero(zeros[3] + 1, z1z16[3], y1y16[3]); - - __syncthreads(); - - int k = offset_k; - int lk = 0; - - while (k < end_k) { - if (k == nextgroup) { - group++; - nextgroup += groupsize; - b_gptq_qzeros_.item4(zeros, group, n); - b_gptq_scales_.item4_h2(scales, group, n); - dequant_4bit_8_prep_zero(zeros[0] + 1, z1z16[0], y1y16[0]); - dequant_4bit_8_prep_zero(zeros[1] + 1, z1z16[1], y1y16[1]); - dequant_4bit_8_prep_zero(zeros[2] + 1, z1z16[2], y1y16[2]); - dequant_4bit_8_prep_zero(zeros[3] + 1, z1z16[3], y1y16[3]); - } - - for (int p = 0; p < 4; p++) { - half2 dq[4][4]; - const int4* b_ptr4 = (int4*)b_ptr; - int4 load_int4 = *b_ptr4; - - dequant_4bit_8_gptq(load_int4.x, dq[0], z1z16[0], y1y16[0], size_n, false); - dequant_4bit_8_gptq(load_int4.y, dq[1], z1z16[1], y1y16[1], size_n, false); - dequant_4bit_8_gptq(load_int4.z, dq[2], z1z16[2], y1y16[2], size_n, false); - dequant_4bit_8_gptq(load_int4.w, dq[3], z1z16[3], y1y16[3], size_n, false); - - b_ptr += size_n; - // half* dqh = (half*)dq; - if (b_q_perm) { - for (int j = 0; j < 4; j++) { - for (int v = 0; v < 4; v++) - dq[v][j] = __hmul2(scales[v], dq[v][j]); - b_.set4( - perm[lk++], n, __low2half(dq[0][j]), __low2half(dq[1][j]), __low2half(dq[2][j]), __low2half(dq[3][j])); - b_.set4( - perm[lk++], - n, - __high2half(dq[0][j]), - __high2half(dq[1][j]), - __high2half(dq[2][j]), - __high2half(dq[3][j])); - } - } else { - for (int j = 0; j < 4; j++) { - for (int v = 0; v < 4; v++) - dq[v][j] = __hmul2(scales[v], dq[v][j]); - b_.set4( - offset_k + lk++, - n, - __low2half(dq[0][j]), - __low2half(dq[1][j]), - __low2half(dq[2][j]), - __low2half(dq[3][j])); - b_.set4( - offset_k + lk++, - n, - __high2half(dq[0][j]), - __high2half(dq[1][j]), - __high2half(dq[2][j]), - __high2half(dq[3][j])); - } - } - } - k += 32; - } -} - -__global__ void reconstruct_exllama_3bit_kernel( - const uint32_t* __restrict__ b_q_weight, - const int* __restrict__ b_q_perm, - const uint32_t* __restrict__ b_gptq_qzeros, - const half* __restrict__ b_gptq_scales, - const int size_k, - const int size_n, - const int groups, - half* __restrict__ b) { - MatrixView_half_rw b_(b, size_k, size_n); - MatrixView_q3_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n); - MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n); - - auto offset_k = BLOCK_KN_SIZE * blockIdx.y; - auto offset_n = BLOCK_KN_SIZE * blockIdx.x * 4; - - int end_k = min(offset_k + BLOCK_KN_SIZE, size_k); - - // Preload remapping table - __shared__ int perm[BLOCK_KN_SIZE]; - auto t = threadIdx.x; - - if (b_q_perm) { - if (offset_k + t < size_k) perm[t] = b_q_perm[offset_k + t]; - } - - // Column - int n = offset_n + t * 4; - if (n >= size_n) return; - - // Find initial group - int groupsize = size_k / groups; - int group = offset_k / groupsize; - int nextgroup = offset_k + groupsize; - - // b offset - int qk = offset_k / 32 * 3; - - const uint32_t* b_ptr = b_q_weight + qk * size_n + n; - - // Initial zeros/scale - int zeros[4]; - half2 scales[4]; - b_gptq_qzeros_.item4(zeros, group, n); - b_gptq_scales_.item4_h2(scales, group, n); - - __syncthreads(); - - int k = offset_k; - int lk = 0; - - while (k < end_k) { - if (k == nextgroup) { - group++; - nextgroup += groupsize; - b_gptq_qzeros_.item4(zeros, group, n); - b_gptq_scales_.item4_h2(scales, group, n); - } - - for (int p = 0; p < 1; p++) { - int4 load_int4[3]; - load_int4[0] = *((int4*)b_ptr); - b_ptr += size_n; - load_int4[1] = *((int4*)b_ptr); - b_ptr += size_n; - load_int4[2] = *((int4*)b_ptr); - b_ptr += size_n; - - half2 dq[4][16]; - dequant_3bit_32(load_int4[0].x, load_int4[1].x, load_int4[2].x, dq[0], size_n, zeros[0] + 1); - dequant_3bit_32(load_int4[0].y, load_int4[1].y, load_int4[2].y, dq[1], size_n, zeros[1] + 1); - dequant_3bit_32(load_int4[0].z, load_int4[1].z, load_int4[2].z, dq[2], size_n, zeros[2] + 1); - dequant_3bit_32(load_int4[0].w, load_int4[1].w, load_int4[2].w, dq[3], size_n, zeros[3] + 1); - - if (b_q_perm) { - for (int j = 0; j < 16; j++) { - for (int v = 0; v < 4; v++) - dq[v][j] = __hmul2(scales[v], dq[v][j]); - b_.set4( - perm[lk++], n, __low2half(dq[0][j]), __low2half(dq[1][j]), __low2half(dq[2][j]), __low2half(dq[3][j])); - b_.set4( - perm[lk++], - n, - __high2half(dq[0][j]), - __high2half(dq[1][j]), - __high2half(dq[2][j]), - __high2half(dq[3][j])); - } - } else { - for (int j = 0; j < 16; j++) { - for (int v = 0; v < 4; v++) - dq[v][j] = __hmul2(scales[v], dq[v][j]); - b_.set4( - offset_k + lk++, - n, - __low2half(dq[0][j]), - __low2half(dq[1][j]), - __low2half(dq[2][j]), - __low2half(dq[3][j])); - b_.set4( - offset_k + lk++, - n, - __high2half(dq[0][j]), - __high2half(dq[1][j]), - __high2half(dq[2][j]), - __high2half(dq[3][j])); - } - } - } - k += 32; - } -} - -__global__ void reconstruct_exllama_2bit_kernel( - const uint32_t* __restrict__ b_q_weight, - const int* __restrict__ b_q_perm, - const uint32_t* __restrict__ b_gptq_qzeros, - const half* __restrict__ b_gptq_scales, - const int size_k, - const int size_n, - const int groups, - half* __restrict__ b) { - MatrixView_half_rw b_(b, size_k, size_n); - MatrixView_q2_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n); - MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n); - - auto offset_k = BLOCK_KN_SIZE * blockIdx.y; - auto offset_n = BLOCK_KN_SIZE * blockIdx.x * 4; - - int end_k = min(offset_k + BLOCK_KN_SIZE, size_k); - - // Preload remapping table - __shared__ int perm[BLOCK_KN_SIZE]; - auto t = threadIdx.x; - - if (b_q_perm) { - if (offset_k + t < size_k) perm[t] = b_q_perm[offset_k + t]; - } - - // Column - int n = offset_n + t * 4; - if (n >= size_n) return; - - // Find initial group - int groupsize = size_k / groups; - int group = offset_k / groupsize; - int nextgroup = offset_k + groupsize; - - // b offset - int qk = offset_k / (32 / 2); - - const uint32_t* b_ptr = b_q_weight + qk * size_n + n; - - // Initial zeros/scale - int zeros[4]; - half2 scales[4]; - b_gptq_qzeros_.item4(zeros, group, n); - b_gptq_scales_.item4_h2(scales, group, n); - - __syncthreads(); - - int k = offset_k; - int lk = 0; - - while (k < end_k) { - if (k == nextgroup) { - group++; - nextgroup += groupsize; - b_gptq_qzeros_.item4(zeros, group, n); - b_gptq_scales_.item4_h2(scales, group, n); - } - - for (int p = 0; p < 2; p++) { - const int4* b_ptr4 = (int4*)b_ptr; - int4 load_int4 = *b_ptr4; - - half2 dq[4][8]; - dequant_2bit_16(load_int4.x, dq[0], size_n, zeros[0] + 1); - dequant_2bit_16(load_int4.y, dq[1], size_n, zeros[1] + 1); - dequant_2bit_16(load_int4.z, dq[2], size_n, zeros[2] + 1); - dequant_2bit_16(load_int4.w, dq[3], size_n, zeros[3] + 1); - - b_ptr += size_n; - // half* dqh = (half*)dq; - if (b_q_perm) { - for (int j = 0; j < 8; j++) { - for (int v = 0; v < 4; v++) - dq[v][j] = __hmul2(scales[v], dq[v][j]); - b_.set4( - perm[lk++], n, __low2half(dq[0][j]), __low2half(dq[1][j]), __low2half(dq[2][j]), __low2half(dq[3][j])); - b_.set4( - perm[lk++], - n, - __high2half(dq[0][j]), - __high2half(dq[1][j]), - __high2half(dq[2][j]), - __high2half(dq[3][j])); - } - } else { - for (int j = 0; j < 8; j++) { - for (int v = 0; v < 4; v++) - dq[v][j] = __hmul2(scales[v], dq[v][j]); - b_.set4( - offset_k + lk++, - n, - __low2half(dq[0][j]), - __low2half(dq[1][j]), - __low2half(dq[2][j]), - __low2half(dq[3][j])); - b_.set4( - offset_k + lk++, - n, - __high2half(dq[0][j]), - __high2half(dq[1][j]), - __high2half(dq[2][j]), - __high2half(dq[3][j])); - } - } - } - k += 32; - } -} - -void reconstruct_exllama( - const uint32_t* b_q_weight, - const uint32_t* b_gptq_qzeros, - const half* b_gptq_scales, - const int* b_q_perm, - half* out, - int height, - int width, - int groups, - int bit) { - dim3 blockDim, gridDim; - blockDim.x = BLOCK_KN_SIZE; - blockDim.y = 1; - gridDim.y = DIVIDE(height, BLOCK_KN_SIZE); - gridDim.x = DIVIDE(width, BLOCK_KN_SIZE); - - auto reconstruct_exllama_kernel = reconstruct_exllama_4bit_kernel; - if (bit == 2) { - reconstruct_exllama_kernel = reconstruct_exllama_2bit_kernel; - } else if (bit == 3) { - reconstruct_exllama_kernel = reconstruct_exllama_3bit_kernel; - } else if (bit == 8) { - reconstruct_exllama_kernel = reconstruct_exllama_8bit_kernel; - } - - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - reconstruct_exllama_kernel<<>>( - b_q_weight, b_q_perm, b_gptq_qzeros, b_gptq_scales, height, width, groups, out); -} - -__global__ void gemm_half_q_half_alt_4bit_kernel( - const half2* __restrict__ vec, - const uint32_t* __restrict__ mat, - half* __restrict__ mul, - const half* __restrict__ scales, - const uint32_t* __restrict__ zeros, - const int* __restrict__ g_idx, - int batch, - int height, - int width) { - int zero_width = width / 8; - int vec_height = height * 4; - const int blockwidth2 = BLOCK_KN_SIZE / 2; - auto b = blockIdx.y * BLOCK_M_SIZE_MAX; - int b_end = min(BLOCK_M_SIZE_MAX, batch - b); - auto h = BLOCK_KN_SIZE * blockIdx.z / 8; - int h_end = min(BLOCK_KN_SIZE / 8, height - h) * 4; - auto w = BLOCK_KN_SIZE * blockIdx.x + threadIdx.x; - - __shared__ half2 blockvec[BLOCK_M_SIZE_MAX][blockwidth2]; - if (threadIdx.x < h_end) { - for (int m = 0; m < b_end; ++m) { - blockvec[m][threadIdx.x] = vec[(m + b) * vec_height + blockIdx.z * BLOCK_KN_SIZE / 2 + threadIdx.x]; - } - } - - __shared__ half2 deq2[256][8]; - auto val = threadIdx.x / 8; - auto off = threadIdx.x % 8; - for (; val < 256; val += BLOCK_KN_SIZE / 8) { - deq2[val][off] = __halves2half2(__int2half_rn(val & 0xF), __int2half_rn(val >> 4)); - } - - if (blockIdx.z == 0) { - for (int m = 0; m < b_end; m++) - mul[(b + m) * width + w] = __int2half_rn(0); - } - __syncthreads(); - - int i = width * h + w; - int g_h = h * 8; - int k = 0; - int z_w = w / 8; - int z_mod = (w % 8) * 4; - half2 res2; - half res[BLOCK_M_SIZE_MAX] = {}; - - unsigned int tmp; - while (k < h_end) { - tmp = mat[i]; - half2 scales_tmp[4]; - half2 zeros_tmp[4]; - for (int tmp_k = 0; tmp_k < 4; tmp_k++) { - int g = g_idx[g_h + (k + tmp_k) * 2]; - int g2 = g_idx[g_h + (k + tmp_k) * 2 + 1]; - half scale_f = scales[g * width + w]; - half scale_f2 = scales[g2 * width + w]; - half2 scale = __halves2half2(scale_f, scale_f2); - half2 zero = __halves2half2( - __hmul(scale_f, __int2half_rn(-((zeros[g * zero_width + z_w] >> z_mod) & 0xF) - 1)), - __hmul(scale_f2, __int2half_rn(-((zeros[g2 * zero_width + z_w] >> z_mod) & 0xF) - 1))); - scales_tmp[tmp_k] = scale; - zeros_tmp[tmp_k] = zero; - } - for (int m = 0; m < b_end; m++) { -#ifndef USE_ROCM - res2 = {}; -#else - res2.x = __half_as_ushort(__float2half(0)); - res2.y = __half_as_ushort(__float2half(0)); -#endif - res2 = __hfma2(__hfma2(deq2[(tmp >> 0) & 0xff][off], scales_tmp[0], zeros_tmp[0]), blockvec[m][k + 0], res2); - res2 = __hfma2(__hfma2(deq2[(tmp >> 8) & 0xff][off], scales_tmp[1], zeros_tmp[1]), blockvec[m][k + 1], res2); - res2 = __hfma2(__hfma2(deq2[(tmp >> 16) & 0xff][off], scales_tmp[2], zeros_tmp[2]), blockvec[m][k + 2], res2); - res2 = __hfma2(__hfma2(deq2[(tmp >> 24) & 0xff][off], scales_tmp[3], zeros_tmp[3]), blockvec[m][k + 3], res2); -#ifndef USE_ROCM - res[m] = __hadd(res[m], __hadd(res2.x, res2.y)); -#else - res[m] = __hadd(res[m], __hadd(__ushort_as_half(res2.x), __ushort_as_half(res2.y))); -#endif - } - i += width; - k += 4; - } - for (int m = 0; m < b_end; m++) { - atomicAdd(&mul[(b + m) * width + w], res[m]); - } -} - -__global__ void gemm_half_q_half_alt_8bit_kernel( - const half2* __restrict__ vec, - const uint32_t* __restrict__ mat, - half* __restrict__ mul, - const half* __restrict__ scales, - const uint32_t* __restrict__ zeros, - const int* __restrict__ g_idx, - int batch, - int height, - int width) { - int zero_width = width / 4; - int vec_height = height * 2; - const int blockwidth2 = BLOCK_KN_SIZE / 2; - auto b = blockIdx.y * BLOCK_M_SIZE_MAX; - int b_end = min(BLOCK_M_SIZE_MAX, batch - b); - auto h = BLOCK_KN_SIZE * blockIdx.z / 4; - int h_end = min(BLOCK_KN_SIZE / 4, height - h) * 2; - auto w = BLOCK_KN_SIZE * blockIdx.x + threadIdx.x; - - __shared__ half2 blockvec[BLOCK_M_SIZE_MAX][blockwidth2]; - if (threadIdx.x < h_end) { - for (int m = 0; m < b_end; ++m) { - blockvec[m][threadIdx.x] = vec[(m + b) * vec_height + blockIdx.z * BLOCK_KN_SIZE / 2 + threadIdx.x]; - } - } - - if (blockIdx.z == 0) { - for (int m = 0; m < b_end; m++) - mul[(b + m) * width + w] = __int2half_rn(0); - } - __syncthreads(); - - int i = width * h + w; - int g_h = h * 4; - int k = 0; - int z_w = w / 4; - int z_mod = (w % 4) * 8; - half2 res2; - half res[BLOCK_M_SIZE_MAX] = {}; - - unsigned int tmp; - while (k < h_end) { - tmp = mat[i]; - half2 scales_tmp[2]; - half2 zeros_tmp[2]; - for (int tmp_k = 0; tmp_k < 2; tmp_k++) { - int g = g_idx[g_h + (k + tmp_k) * 2]; - int g2 = g_idx[g_h + (k + tmp_k) * 2 + 1]; - half scale_f = scales[g * width + w]; - half scale_f2 = scales[g2 * width + w]; - half2 scale = __halves2half2(scale_f, scale_f2); - half2 zero = __halves2half2( - __hmul(scale_f, __int2half_rn(-((zeros[g * zero_width + z_w] >> z_mod) & 0xff) - 1)), - __hmul(scale_f2, __int2half_rn(-((zeros[g2 * zero_width + z_w] >> z_mod) & 0xff) - 1))); - scales_tmp[tmp_k] = scale; - zeros_tmp[tmp_k] = zero; - } - for (int m = 0; m < b_end; m++) { -#ifndef USE_ROCM - res2 = {}; -#else - res2.x = __half_as_ushort(__float2half(0)); - res2.y = __half_as_ushort(__float2half(0)); -#endif - half2 v12 = __halves2half2(__int2half_rn(tmp & 0xFF), __int2half_rn((tmp >> 8) & 0xFF)); - res2 = __hfma2(__hfma2(v12, scales_tmp[0], zeros_tmp[0]), blockvec[m][k + 0], res2); - half2 v34 = __halves2half2(__int2half_rn((tmp >> 16) & 0xFF), __int2half_rn((tmp >> 24) & 0xFF)); - res2 = __hfma2(__hfma2(v34, scales_tmp[1], zeros_tmp[1]), blockvec[m][k + 1], res2); -#ifndef USE_ROCM - res[m] = __hadd(res[m], __hadd(res2.x, res2.y)); -#else - res[m] = __hadd(res[m], __hadd(__ushort_as_half(res2.x), __ushort_as_half(res2.y))); -#endif - } - i += width; - k += 2; - } - for (int m = 0; m < b_end; m++) { - atomicAdd(&mul[(b + m) * width + w], res[m]); - } -} - -void gemm_half_q_half_alt( - const half* a, - const uint32_t* b_q_weight, - const uint32_t* b_gptq_qzeros, - const half* b_gptq_scales, - const int* b_g_idx, - half* c, - int size_m, - int size_n, - int size_k, - int bit) { - dim3 blockDim, gridDim; - blockDim.x = BLOCK_KN_SIZE; - blockDim.y = 1; - blockDim.z = 1; - gridDim.x = DIVIDE(size_n, BLOCK_KN_SIZE); - gridDim.y = DIVIDE(size_m, BLOCK_M_SIZE_MAX); - gridDim.z = DIVIDE(size_k, BLOCK_KN_SIZE); - - auto kernel = gemm_half_q_half_alt_4bit_kernel; - if (bit == 8) { - kernel = gemm_half_q_half_alt_8bit_kernel; - } - - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - kernel<<>>( - (const half2*)a, b_q_weight, c, b_gptq_scales, b_gptq_qzeros, b_g_idx, size_m, size_k / 32 * bit, size_n); -} - -template -__global__ void reconstruct_gptq_kernel( - const uint32_t* __restrict__ w, - const half* __restrict__ w_scales, - const uint32_t* __restrict__ w_zeros, - const int* __restrict__ g_idx, - const int height, - const int width, - const int group, - half* __restrict__ out) { - // Start of block - - auto column = BLOCK_KN_SIZE * blockIdx.x + threadIdx.x; - auto row = blockIdx.y * 32 / bit; - if (column >= width) return; - - // Views - - MatrixView_half_rw out_(out, height, width); - MatrixView_half w_scales_(w_scales, group, width); - T w_zeros_(w_zeros, group, width); - - uint32_t w_read = w[blockIdx.y * width + column]; - half* out_ptr = out_.item_ptr(row, column); - -#pragma unroll - for (int s = 0; s < 32; s += bit) { - int group = g_idx[row + s / bit]; - half w_scale = w_scales_.item(group, column); - uint32_t w_zero = w_zeros_.item(group, column) + 1; - half w_item = __hmul(__int2half_rn((int)((w_read >> s) & ((1 << bit) - 1)) - w_zero), w_scale); - *out_ptr = w_item; - out_ptr += out_.width; - } -} - -__global__ void reconstruct_gptq_3bit_kernel( - const uint32_t* __restrict__ w, - const half* __restrict__ w_scales, - const uint32_t* __restrict__ w_zeros, - const int* __restrict__ g_idx, - const int height, - const int width, - const int group, - half* __restrict__ out) { - // Start of block - auto column = BLOCK_KN_SIZE * blockIdx.x + threadIdx.x; - auto row = blockIdx.y * 32; - if (column >= width) return; - - // Views - - MatrixView_half_rw out_(out, height, width); - MatrixView_half w_scales_(w_scales, group, width); - MatrixView_q3_row w_zeros_(w_zeros, group, width); - - uint32_t w1 = w[(blockIdx.y * 3) * width + column]; - uint32_t w2 = w[(blockIdx.y * 3 + 1) * width + column]; - uint32_t w3 = w[(blockIdx.y * 3 + 2) * width + column]; - half* out_ptr = out_.item_ptr(row, column); - -#pragma unroll - for (int i = 0; i < 32; i += 1) { - int group = g_idx[row + i]; - half w_scale = w_scales_.item(group, column); - uint32_t w_zero = w_zeros_.item(group, column) + 1; - int w_item; - if (i == 10) { - w_item = (w1 >> 30) | ((w2 << 2) & 0x4); - } else if (i == 21) { - w_item = (w2 >> 31) | ((w3 << 1) & 0x6); - } else if (i < 10) { - w_item = ((w1 >> (i * 3)) & 0x7); - } else if (i < 21) { - w_item = ((w2 >> (i * 3 - 32)) & 0x7); - } else { - w_item = ((w3 >> (i * 3 - 64)) & 0x7); - } - *out_ptr = __hmul(__int2half_rn(w_item - w_zero), w_scale); - out_ptr += out_.width; - } -} - -void reconstruct_gptq( - const uint32_t* b_q_weight, - const uint32_t* b_gptq_qzeros, - const half* b_gptq_scales, - const int* b_g_idx, - half* out, - int height, - int width, - int groups, - int bit) { - dim3 blockDim, gridDim; - blockDim.x = BLOCK_KN_SIZE; - blockDim.y = 1; - gridDim.y = DIVIDE(height, 32 / bit); - gridDim.x = DIVIDE(width, BLOCK_KN_SIZE); - - auto kernel = reconstruct_gptq_kernel; - if (bit == 2) { - kernel = reconstruct_gptq_kernel; - } else if (bit == 8) { - kernel = reconstruct_gptq_kernel; - } else if (bit == 3) { - kernel = reconstruct_gptq_3bit_kernel; - gridDim.y = DIVIDE(height, 32); - } - - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - kernel<<>>( - b_q_weight, b_gptq_scales, b_gptq_qzeros, b_g_idx, height, width, groups, out); -} - -void gemm_half_q_half_cuda( - cublasHandle_t cublas_handle, - const half* a, - const uint32_t* b_q_weight, - const uint32_t* b_gptq_qzeros, - const half* b_gptq_scales, - const int* b_g_idx, - half* c, - half* temp_dq, - int size_m, - int size_n, - int size_k, - int groups, - bool use_shuffle, - int bit) { - bool use_reconstruct; - if (use_shuffle) { - use_reconstruct = ((bit == 8 && size_m > MAX_Q_GEMM_ROWS_8BIT) || (bit != 8 && size_m > MAX_Q_GEMM_ROWS)); - } else { - // The 2/3-bit kernels are somehow slower than dequant + gemm baseline, so - // we disabled them for now. - use_reconstruct = (bit < 4 || size_m > MAX_ALT_GEMM_ROWS); - } - if (use_reconstruct) { - // Reconstruct FP16 matrix, then cuBLAS - if (use_shuffle) { - reconstruct_exllama(b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, temp_dq, size_k, size_n, groups, bit); - } else { - reconstruct_gptq(b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, temp_dq, size_k, size_n, groups, bit); - } - - const half alpha = __float2half(1.0f); - const half beta = __float2half(0.0f); - cublasHgemm( - cublas_handle, - CUBLAS_OP_N, - CUBLAS_OP_N, - size_n, - size_m, - size_k, - &alpha, - temp_dq, - size_n, - a, - size_k, - &beta, - c, - size_n); - } else if (use_shuffle) { - // Quantized matmul - int max_chunks = size_m / BLOCK_M_SIZE_MAX; - int last_chunk = max_chunks * BLOCK_M_SIZE_MAX; - int last_chunk_size = size_m - last_chunk; - - if (max_chunks) { - gemm_half_q_half_cuda_part( - a, - b_q_weight, - b_gptq_qzeros, - b_gptq_scales, - b_g_idx, - c, - last_chunk, - size_n, - size_k, - BLOCK_M_SIZE_MAX, - groups, - bit); - } - - if (last_chunk_size) { - gemm_half_q_half_cuda_part( - a + last_chunk * size_k, - b_q_weight, - b_gptq_qzeros, - b_gptq_scales, - b_g_idx, - c + last_chunk * size_n, - last_chunk_size, - size_n, - size_k, - last_chunk_size, - groups, - bit); - } - } else { - gemm_half_q_half_alt(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, c, size_m, size_n, size_k, bit); - } -} - -__global__ void shuffle_4bit_kernel(uint32_t* __restrict__ b_q_weight, const int size_k, const int size_n) { - auto n = blockIdx.x * THREADS_X + threadIdx.x; - if (n >= size_n) return; - int k = 0; - uint32_t* b_ptr = b_q_weight + n; - while (k < size_k) { - shuffle_4bit_8(b_ptr, size_n); - b_ptr += 1 * size_n; - k += 8; - } -} - -__global__ void shuffle_8bit_kernel(uint32_t* __restrict__ b_q_weight, const int size_k, const int size_n) { - auto n = blockIdx.x * THREADS_X + threadIdx.x; - if (n >= size_n) return; - int k = 0; - uint32_t* b_ptr = b_q_weight + n; - while (k < size_k) { - shuffle_8bit_4(b_ptr, size_n); - b_ptr += 1 * size_n; - k += 4; - } -} - -__global__ void shuffle_2bit_kernel(uint32_t* __restrict__ b_q_weight, const int size_k, const int size_n) { - auto n = blockIdx.x * THREADS_X + threadIdx.x; - if (n >= size_n) return; - int k = 0; - uint32_t* b_ptr = b_q_weight + n; - while (k < size_k) { - shuffle_2bit_16(b_ptr, size_n); - b_ptr += 1 * size_n; - k += 16; - } -} - -__global__ void shuffle_3bit_kernel(uint32_t* __restrict__ b_q_weight, const int size_k, const int size_n) { - auto n = blockIdx.x * THREADS_X + threadIdx.x; - if (n >= size_n) return; - int k = 0; - uint32_t* b_ptr = b_q_weight + n; - while (k < size_k) { - shuffle_3bit_32(b_ptr, size_n); - b_ptr += 3 * size_n; - k += 32; - } -} - -__global__ void make_sequential_4bit_kernel( - const uint32_t* __restrict__ w, uint32_t* __restrict__ w_new, const int* __restrict__ q_perm, const int w_width) { - const uint64_t* w2 = (uint64_t*)w; - uint64_t* w_new2 = (uint64_t*)w_new; - int w2_stride = w_width >> 1; - auto w2_column = THREADS_X * blockIdx.x + threadIdx.x; - if (w2_column >= w2_stride) return; - auto w_new2_row = blockIdx.y; - int q_perm_idx = w_new2_row << 3; - uint64_t dst = 0; - -#pragma unroll - for (int i = 0; i < 8; i++) { - int source_row = q_perm[q_perm_idx++]; - - int w2_row = source_row >> 3; - int w2_subrow = source_row & 0x07; - int w2_row_shift = w2_subrow << 2; - int wnew2_row_shift = i << 2; - - uint64_t src = w2[w2_row * w2_stride + w2_column]; - src >>= w2_row_shift; - src &= 0x0000000f0000000f; - src <<= wnew2_row_shift; - dst |= src; - } - w_new2[w_new2_row * w2_stride + w2_column] = dst; -} - -__global__ void make_sequential_2bit_kernel( - const uint32_t* __restrict__ w, uint32_t* __restrict__ w_new, const int* __restrict__ q_perm, const int w_width) { - const uint64_t* w2 = (uint64_t*)w; - uint64_t* w_new2 = (uint64_t*)w_new; - int w2_stride = w_width >> 1; - auto w2_column = THREADS_X * blockIdx.x + threadIdx.x; - if (w2_column >= w2_stride) return; - auto w_new2_row = blockIdx.y; - int q_perm_idx = w_new2_row << 4; - uint64_t dst = 0; - -#pragma unroll - for (int i = 0; i < 16; i++) { - int source_row = q_perm[q_perm_idx++]; - - int w2_row = source_row >> 4; - int w2_subrow = source_row & 0x0f; - int w2_row_shift = w2_subrow << 1; - int wnew2_row_shift = i << 1; - - uint64_t src = w2[w2_row * w2_stride + w2_column]; - src >>= w2_row_shift; - src &= 0x0000000300000003; - src <<= wnew2_row_shift; - dst |= src; - } - w_new2[w_new2_row * w2_stride + w2_column] = dst; -} - -__global__ void make_sequential_3bit_kernel( - const uint32_t* __restrict__ w, uint32_t* __restrict__ w_new, const int* __restrict__ q_perm, const int w_width) { - auto w_column = THREADS_X * blockIdx.x + threadIdx.x; - if (w_column >= w_width) return; - auto w_new_row = blockIdx.y * 3; - auto q_perm_idx = blockIdx.y << 5; - uint32_t dst[3] = {0, 0, 0}; - -#pragma unroll - for (int i = 0; i < 32; i++) { - int source_row = q_perm[q_perm_idx++]; - int z_w = (source_row / 32) * 3; - int z_mod = source_row % 32; - int z_bit; - - if (z_mod != 10) { - if (z_mod != 21) { - z_bit = z_mod; - if (z_bit > 21) { - z_bit *= 3; - z_bit -= 64; - z_w += 2; - } else if (z_bit > 10) { - z_bit *= 3; - z_bit -= 32; - z_w += 1; - } else { - z_bit *= 3; - } - } else { - z_w += 1; - } - } - - uint64_t src; - if (z_mod == 10) { - src = (w[z_w * w_width + w_column] >> 30) | ((w[(z_w + 1) * w_width + w_column] << 2) & 0x4); - } else if (z_mod == 21) { - src = (w[z_w * w_width + w_column] >> 31) | ((w[(z_w + 1) * w_width + w_column] << 1) & 0x6); - } else { - src = w[z_w * w_width + w_column]; - src >>= z_bit; - src &= 0x07; - } - - z_w = 0; - if (i != 10) { - if (i != 21) { - z_bit = i; - if (z_bit > 21) { - z_bit *= 3; - z_bit -= 64; - z_w += 2; - } else if (z_bit > 10) { - z_bit *= 3; - z_bit -= 32; - z_w += 1; - } else { - z_bit *= 3; - } - } else { - z_w += 1; - } - } - if (i == 10) { - dst[z_w] |= (src & 0x03) << 30; - dst[z_w + 1] |= ((src & 0x4) >> 2); - } else if (i == 21) { - dst[z_w] |= (src & 0x01) << 31; - dst[z_w + 1] |= ((src & 0x6) >> 1); - } else { - dst[z_w] |= (src << z_bit); - } - } - w_new[w_new_row * w_width + w_column] = dst[0]; - w_new[(w_new_row + 1) * w_width + w_column] = dst[1]; - w_new[(w_new_row + 2) * w_width + w_column] = dst[2]; -} - -__global__ void make_sequential_8bit_kernel( - const uint32_t* __restrict__ w, uint32_t* __restrict__ w_new, const int* __restrict__ q_perm, const int w_width) { - const uint64_t* w2 = (uint64_t*)w; - uint64_t* w_new2 = (uint64_t*)w_new; - int w2_stride = w_width >> 1; - auto w2_column = THREADS_X * blockIdx.x + threadIdx.x; - if (w2_column >= w2_stride) return; - auto w_new2_row = blockIdx.y; - int q_perm_idx = w_new2_row << 2; - uint64_t dst = 0; - -#pragma unroll - for (int i = 0; i < 4; i++) { - int source_row = q_perm[q_perm_idx++]; - - int w2_row = source_row >> 2; - int w2_subrow = source_row & 0x03; - int w2_row_shift = w2_subrow << 3; - int wnew2_row_shift = i << 3; - - uint64_t src = w2[w2_row * w2_stride + w2_column]; - src >>= w2_row_shift; - src &= 0x000000ff000000ff; - src <<= wnew2_row_shift; - dst |= src; - } - w_new2[w_new2_row * w2_stride + w2_column] = dst; -} - -void shuffle_exllama_weight(uint32_t* q_weight, int* q_perm, int height, int width, int bit) { - if (q_perm) { - uint32_t* new_qweight = NULL; - cudaMalloc(&new_qweight, height / 32 * bit * width * sizeof(uint32_t)); - - dim3 blockDim, gridDim; - blockDim.x = THREADS_X; - blockDim.y = 1; - gridDim.x = DIVIDE(width, THREADS_X); - gridDim.y = height / 32 * bit; - - auto kernel = make_sequential_4bit_kernel; - if (bit == 2) { - kernel = make_sequential_2bit_kernel; - } else if (bit == 3) { - kernel = make_sequential_3bit_kernel; - gridDim.y = height / 32; - } else if (bit == 8) { - kernel = make_sequential_8bit_kernel; - } - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - kernel<<>>(q_weight, new_qweight, q_perm, width); - // Replace qweights - cudaMemcpyAsync(q_weight, new_qweight, height / 32 * bit * width * sizeof(uint32_t), cudaMemcpyDeviceToDevice); - // Cleanup - cudaDeviceSynchronize(); - cudaFree(new_qweight); - } - dim3 blockDim, gridDim; - blockDim.x = THREADS_X; - blockDim.y = 1; - gridDim.x = DIVIDE(width, THREADS_X); - gridDim.y = 1; - auto shuffle_kernel = shuffle_4bit_kernel; - if (bit == 2) { - shuffle_kernel = shuffle_2bit_kernel; - } else if (bit == 3) { - shuffle_kernel = shuffle_3bit_kernel; - } else if (bit == 8) { - shuffle_kernel = shuffle_8bit_kernel; - } - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - shuffle_kernel<<>>(q_weight, height, width); -} - -} // namespace gptq -} // namespace sglang - -torch::Tensor gptq_gemm( - torch::Tensor a, - torch::Tensor b_q_weight, - torch::Tensor b_gptq_qzeros, - torch::Tensor b_gptq_scales, - torch::Tensor b_g_idx, - bool use_shuffle, - int64_t bit) { - const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); - auto options = torch::TensorOptions().dtype(a.dtype()).device(a.device()); - at::Tensor c = torch::empty({a.size(0), b_q_weight.size(1)}, options); - at::Tensor temp_dq = torch::empty({b_q_weight.size(0) * 32 / bit, b_q_weight.size(1)}, options); - - sglang::gptq::gemm_half_q_half_cuda( - at::cuda::getCurrentCUDABlasHandle(), - (const half*)a.data_ptr(), - (const uint32_t*)b_q_weight.data_ptr(), - (const uint32_t*)b_gptq_qzeros.data_ptr(), - (const half*)b_gptq_scales.data_ptr(), - b_g_idx.device().is_meta() ? NULL : (const int*)b_g_idx.data_ptr(), - (half*)c.data_ptr(), - (half*)temp_dq.data_ptr(), - c.size(0), // m - c.size(1), // n - a.size(1), // k - b_gptq_qzeros.size(0), // group number - use_shuffle, - bit); - return c; -} - -void gptq_shuffle(torch::Tensor q_weight, torch::Tensor q_perm, int64_t bit) { - const at::cuda::OptionalCUDAGuard device_guard(device_of(q_weight)); - sglang::gptq::shuffle_exllama_weight( - (uint32_t*)q_weight.data_ptr(), - q_perm.device().is_meta() || q_perm.numel() == 0 ? NULL : (int*)q_perm.data_ptr(), - q_weight.size(0) * 32 / bit, - q_weight.size(1), - bit); -} diff --git a/python/sglang/kernels/aot/csrc/gemm/gptq/matrix_view.cuh b/python/sglang/kernels/aot/csrc/gemm/gptq/matrix_view.cuh deleted file mode 100644 index 3dfc8794c..000000000 --- a/python/sglang/kernels/aot/csrc/gemm/gptq/matrix_view.cuh +++ /dev/null @@ -1,269 +0,0 @@ -/* -Adapted from https://github.com/turboderp/exllamav2 and -https://github.com/turboderp/exllama -*/ - -#ifndef _matrix_view_cuh -#define _matrix_view_cuh - -#include -#include - -#include "qdq_util.cuh" - -namespace sglang { -namespace gptq { - -class MatrixView_half { - public: - const half* data; - const int height; - const int width; - - __device__ __forceinline__ MatrixView_half(const half* data, const int height, const int width) - : data(data), height(height), width(width) {} - - __device__ __forceinline__ half item(int row, int column) const { - return data[row * width + column]; - } - __device__ __forceinline__ half2 item_half2(int row, int column) const { - return ((half2*)data)[(row * width + column) / 2]; - } - __device__ __forceinline__ half2 item_half2half2(int row, int column) const { - return __half2half2(data[row * width + column]); - } - __device__ __forceinline__ const half* item_ptr(int row, int column) const { - return &data[row * width + column]; - } - - __device__ __forceinline__ void item4(half (&items)[4], int row, int column) const { - half2* ptr = (half2*)item_ptr(row, column); - half2 i01 = ptr[0]; - half2 i23 = ptr[1]; - items[0] = __low2half(i01); - items[1] = __high2half(i01); - items[2] = __low2half(i23); - items[3] = __high2half(i23); - } - __device__ __forceinline__ void item4_f(float (&items)[4], int row, int column) const { - half2* ptr = (half2*)item_ptr(row, column); - half2 i01 = ptr[0]; - half2 i23 = ptr[1]; - items[0] = __half2float(__low2half(i01)); - items[1] = __half2float(__high2half(i01)); - items[2] = __half2float(__low2half(i23)); - items[3] = __half2float(__high2half(i23)); - } - - __device__ __forceinline__ void item4_h2(half2 (&items)[4], int row, int column) const { - half2* ptr = (half2*)item_ptr(row, column); - half2 i01 = ptr[0]; - half2 i23 = ptr[1]; - items[0] = __half2half2(__low2half(i01)); - items[1] = __half2half2(__high2half(i01)); - items[2] = __half2half2(__low2half(i23)); - items[3] = __half2half2(__high2half(i23)); - } -}; - -class MatrixView_half_rw { - public: - half* data; - const int height; - const int width; - - __device__ __forceinline__ MatrixView_half_rw(half* data, const int height, const int width) - : data(data), height(height), width(width) {} - - __device__ __forceinline__ half item(int row, int column) const { - return data[row * width + column]; - } - __device__ __forceinline__ half2 item_half2(int row, int column) const { - return ((half2*)data)[(row * width + column) / 2]; - } - __device__ __forceinline__ half* item_ptr(int row, int column) { - return &data[row * width + column]; - } - __device__ __forceinline__ void set(int row, int column, half value) { - data[row * width + column] = value; - } - __device__ __forceinline__ void set_half2(int row, int column, half2 value) { - ((half2*)data)[(row * width + column) / 2] = value; - } - - __device__ __forceinline__ void set4(int row, int column, half v0, half v1, half v2, half v3) { - half2 v01 = __halves2half2(v0, v1); - half2 v23 = __halves2half2(v2, v3); - half2* ptr = (half2*)item_ptr(row, column); - ptr[0] = v01; - ptr[1] = v23; - } -}; - -class MatrixView_q4_row { - public: - const uint32_t* data; - const int height; - const int width; - - __device__ __forceinline__ MatrixView_q4_row(const uint32_t* data, const int height, const int width) - : data(data), height(height), width(width) {} - - __device__ __forceinline__ int item(int row, int column) const { - int shift = (column & 0x07) * 4; - return (data[row * width / 8 + column / 8] >> shift) & 0x0f; - } - - __device__ __forceinline__ void item2(int (&items)[2], int row, int column) const { - int shift = (column & 0x07) * 4; - uint32_t d = data[row * width / 8 + column / 8] >> shift; - items[0] = d & 0x0f; - items[1] = (d >> 4) & 0x0f; - } - - __device__ __forceinline__ void item4(int (&items)[4], int row, int column) const { - int shift = (column & 0x07) * 4; - uint32_t d = data[row * width / 8 + column / 8] >> shift; - items[0] = d & 0x0f; - items[1] = (d >> 4) & 0x0f; - items[2] = (d >> 8) & 0x0f; - items[3] = (d >> 12) & 0x0f; - } -}; - -class MatrixView_q4_column { - public: - const uint32_t* data; - const int height; - const int width; - - __device__ __forceinline__ MatrixView_q4_column(const uint32_t* data, const int height, const int width) - : data(data), height(height), width(width) {} - - __device__ __forceinline__ int item(int row, int column) const { - int shift = (row & 0x07) * 4; - return (data[row / 8 * width + column] >> shift) & 0x0f; - } - - __device__ __forceinline__ uint32_t item_uint32_t(int row, int column) { - return data[row / 8 * width + column]; - } - __device__ __forceinline__ const uint32_t* item_uint32_ptr(int row, int column) { - return &data[row / 8 * width + column]; - } -}; - -class MatrixView_q2_row { - public: - const uint32_t* data; - const int height; - const int width; - - __device__ __forceinline__ MatrixView_q2_row(const uint32_t* data, const int height, const int width) - : data(data), height(height), width(width) {} - - __device__ __forceinline__ int item(int row, int column) const { - int shift = (column & 0x0f) * 2; - return (data[row * width / 16 + column / 16] >> shift) & 0x03; - } - - __device__ __forceinline__ void item2(int (&items)[2], int row, int column) const { - int shift = (column & 0x0f) * 2; - uint32_t d = data[row * width / 16 + column / 16] >> shift; - items[0] = d & 0x03; - items[1] = (d >> 2) & 0x03; - } - - __device__ __forceinline__ void item4(int (&items)[4], int row, int column) const { - int shift = (column & 0x0f) * 2; - uint32_t d = data[row * width / 16 + column / 16] >> shift; - items[0] = d & 0x03; - items[1] = (d >> 2) & 0x03; - items[2] = (d >> 4) & 0x03; - items[3] = (d >> 6) & 0x03; - } -}; - -class MatrixView_q3_row { - public: - const uint32_t* data; - const int height; - const int width; - - __device__ __forceinline__ MatrixView_q3_row(const uint32_t* data, const int height, const int width) - : data(data), height(height), width(width) {} - - __device__ __forceinline__ int item(int row, int column) const { - int z_w = column * 3 / 32; - int z_mod = column & 0x1f; - - if (z_mod == 10) { - return (data[row * width * 3 / 32 + z_w] >> 30) | ((data[row * width * 3 / 32 + (z_w + 1)] << 2) & 0x4); - } else if (z_mod == 21) { - return (data[row * width * 3 / 32 + z_w] >> 31) | ((data[row * width * 3 / 32 + (z_w + 1)] << 1) & 0x6); - } else if (z_mod < 10) { - return (data[row * width * 3 / 32 + z_w] >> (z_mod * 3)) & 0x07; - } else if (z_mod < 21) { - return (data[row * width * 3 / 32 + z_w] >> (z_mod * 3 - 32)) & 0x07; - } else { - return (data[row * width * 3 / 32 + z_w] >> (z_mod * 3 - 64)) & 0x07; - } - } - - __device__ __forceinline__ void item4(int (&items)[4], int row, int column) const { - int shift = (column & 0x1f); - uint32_t d; - if (shift <= 4) { - d = data[row * width / 32 * 3 + column * 3 / 32] >> (shift * 3); - } else if (shift == 8) { - d = (data[row * width / 32 * 3 + column * 3 / 32] >> 24) | - ((data[row * width / 32 * 3 + column * 3 / 32 + 1] & 0x0f) << 8); - } else if (shift <= 16) { - d = data[row * width / 32 * 3 + column * 3 / 32] >> (shift * 3 - 32); - } else if (shift == 20) { - d = (data[row * width / 32 * 3 + column * 3 / 32] >> 28) | - ((data[row * width / 32 * 3 + column * 3 / 32 + 1] & 0xff) << 4); - } else { - d = data[row * width / 32 * 3 + column * 3 / 32] >> (shift * 3 - 64); - } - items[0] = d & 0x07; - items[1] = (d >> 3) & 0x07; - items[2] = (d >> 6) & 0x07; - items[3] = (d >> 9) & 0x07; - } -}; - -class MatrixView_q8_row { - public: - const uint32_t* data; - const int height; - const int width; - - __device__ __forceinline__ MatrixView_q8_row(const uint32_t* data, const int height, const int width) - : data(data), height(height), width(width) {} - - __device__ __forceinline__ int item(int row, int column) const { - int shift = (column & 0x03) * 8; - return (data[row * width / 4 + column / 4] >> shift) & 0xff; - } - - __device__ __forceinline__ void item2(int (&items)[2], int row, int column) const { - int shift = (column & 0x03) * 8; - uint32_t d = data[row * width / 4 + column / 4] >> shift; - items[0] = d & 0xff; - items[1] = (d >> 8) & 0xff; - } - - __device__ __forceinline__ void item4(int (&items)[4], int row, int column) const { - int shift = (column & 0x03) * 2; - uint32_t d = data[row * width / 4 + column / 4] >> shift; - items[0] = d & 0xff; - items[1] = (d >> 8) & 0xff; - items[2] = (d >> 16) & 0xff; - items[3] = (d >> 24) & 0xff; - } -}; - -} // namespace gptq -} // namespace sglang -#endif diff --git a/python/sglang/kernels/aot/csrc/gemm/gptq/qdq_2.cuh b/python/sglang/kernels/aot/csrc/gemm/gptq/qdq_2.cuh deleted file mode 100644 index 4a75d7b56..000000000 --- a/python/sglang/kernels/aot/csrc/gemm/gptq/qdq_2.cuh +++ /dev/null @@ -1,74 +0,0 @@ -/* -Copied from https://github.com/turboderp/exllamav2 -*/ - -#ifndef _qdq_2_cuh -#define _qdq_2_cuh - -#include "qdq_util.cuh" - -namespace sglang { -namespace gptq { - -// Permutation: -// -// ffddbb99 77553311 eeccaa88 66442200 - -__forceinline__ __device__ void shuffle_2bit_16(uint32_t* q, int stride) { - uint32_t qa = q[0]; - uint32_t qb = 0; - -#pragma unroll - for (int i = 0; i < 8; i++) { - uint32_t qa0 = qa & 0x03; - uint32_t qa1 = (qa & 0x0c) >> 2; - qa >>= 4; - qb |= (qa1 << (i * 2 + 16)); - qb |= (qa0 << (i * 2)); - } - q[0] = qb; -} - -__forceinline__ __device__ void dequant_2bit_16(const uint32_t q_0, half2 (&dq)[8], int stride, const uint32_t zero) { - const uint32_t c0 = 0x64006400; - const half y4_ = __float2half_rn(1.0f / 4.0f); - const half y16_ = __float2half_rn(1.0f / 16.0f); - const half y64_ = __float2half_rn(1.0f / 64.0f); - const half2 y4 = __halves2half2(y4_, y4_); - const half2 y16 = __halves2half2(y16_, y16_); - const half2 y64 = __halves2half2(y64_, y64_); - - const half_uint16 z1_(0xe400 | zero); // half(-1024.0f - zero); - const half z4_ = __hsub(__int2half_rn(-256), __int2half_rn(zero)); - const half z16_ = __hsub(__int2half_rn(-64), __int2half_rn(zero)); - const half z64_ = __hsub(__int2half_rn(-16), __int2half_rn(zero)); - const half2 z1 = __half2half2(z1_.as_half); - const half2 z4 = __half2half2(z4_); - const half2 z16 = __half2half2(z16_); - const half2 z64 = __half2half2(z64_); - - uint32_t qa = q_0; - half2_uint32 q0((qa & 0x00030003) | c0); // half2(q[ 0], q[ 1]) + 1024 - half2_uint32 q1((qa & 0x000c000c) | c0); // half2(q[ 2], q[ 3]) * 4 + 1024 - half2_uint32 q2((qa & 0x00300030) | c0); // half2(q[ 4], q[ 5]) * 16 + 1024 - half2_uint32 q3((qa & 0x00c000c0) | c0); // half2(q[ 6], q[ 7]) * 64 + 1024 - qa >>= 8; - half2_uint32 q4((qa & 0x00030003) | c0); // half2(q[ 8], q[ 8]) + 1024 - half2_uint32 q5((qa & 0x000c000c) | c0); // half2(q[10], q[11]) * 4 + 1024 - half2_uint32 q6((qa & 0x00300030) | c0); // half2(q[12], q[13]) * 16 + 1024 - half2_uint32 q7((qa & 0x00c000c0) | c0); // half2(q[14], q[15]) * 64 + 1024 - - dq[0] = __hadd2(q0.as_half2, z1); - dq[1] = __hfma2(q1.as_half2, y4, z4); - dq[2] = __hfma2(q2.as_half2, y16, z16); - dq[3] = __hfma2(q3.as_half2, y64, z64); - dq[4] = __hadd2(q4.as_half2, z1); - dq[5] = __hfma2(q5.as_half2, y4, z4); - dq[6] = __hfma2(q6.as_half2, y16, z16); - dq[7] = __hfma2(q7.as_half2, y64, z64); -} - -} // namespace gptq -} // namespace sglang - -#endif diff --git a/python/sglang/kernels/aot/csrc/gemm/gptq/qdq_3.cuh b/python/sglang/kernels/aot/csrc/gemm/gptq/qdq_3.cuh deleted file mode 100644 index 5996f342d..000000000 --- a/python/sglang/kernels/aot/csrc/gemm/gptq/qdq_3.cuh +++ /dev/null @@ -1,146 +0,0 @@ -#ifndef _qdq_3_cuh -#define _qdq_3_cuh - -#include "qdq_util.cuh" - -namespace sglang { -namespace gptq { -// Permutation: -// -// v9997775 55333111 u8886664 44222000 (u, v lsb) -// vjjjhhhf ffdddbbb uiiiggge eecccaaa -// vtttrrrp ppnnnlll usssqqqo oommmkkk - -__forceinline__ __device__ void shuffle_3bit_32(uint32_t* q, int stride) { - uint32_t qa = q[0 * stride]; - uint32_t qb = q[1 * stride]; - uint32_t qc = q[2 * stride]; - - // qa: aa999888 77766655 54443332 22111000 - // qb: lkkkjjji iihhhggg fffeeedd dcccbbba - // qc: vvvuuutt tsssrrrq qqpppooo nnnmmmll - - uint32_t qd = qc >> 26; - qc <<= 4; - qc |= qb >> 28; - qb <<= 2; - qb |= qa >> 30; - - // qa: ..999888 77766655 54443332 22111000 - // qb: ..jjjiii hhhgggff feeedddc ccbbbaaa - // qc: ..tttsss rrrqqqpp pooonnnm mmlllkkk - // qd: vvvuuu - - uint32_t za = 0; - uint32_t zb = 0; - uint32_t zc = 0; - - for (int i = 0; i < 5; i++) { - uint32_t t0 = qa & 0x07; - uint32_t t1 = (qa & 0x38) >> 3; - qa >>= 6; - za |= (t0 << (i * 3)); - za |= (t1 << (i * 3 + 16)); - } - for (int i = 0; i < 5; i++) { - uint32_t t0 = qb & 0x07; - uint32_t t1 = (qb & 0x38) >> 3; - qb >>= 6; - zb |= (t0 << (i * 3)); - zb |= (t1 << (i * 3 + 16)); - } - for (int i = 0; i < 5; i++) { - uint32_t t0 = qc & 0x07; - uint32_t t1 = (qc & 0x38) >> 3; - qc >>= 6; - zc |= (t0 << (i * 3)); - zc |= (t1 << (i * 3 + 16)); - } - - // za: 9997775 55333111 8886664 44222000 - // zb: jjjhhhf ffdddbbb iiiggge eecccaaa - // zc: tttrrrp ppnnnlll sssqqqo oommmkkk - // qd: vvvuuu - - za |= ((qd & 0x01) >> 0) << 15; - zb |= ((qd & 0x02) >> 1) << 15; - zc |= ((qd & 0x04) >> 2) << 15; - za |= ((qd & 0x08) >> 3) << 31; - zb |= ((qd & 0x10) >> 4) << 31; - zc |= ((qd & 0x20) >> 5) << 31; - - // za: v9997775 55333111 u8886664 44222000 (u, v lsb) - // zb: vjjjhhhf ffdddbbb uiiiggge eecccaaa - // zc: vtttrrrp ppnnnlll usssqqqo oommmkkk - - q[0 * stride] = za; - q[1 * stride] = zb; - q[2 * stride] = zc; -} - -__forceinline__ __device__ void dequant_3bit_32( - const uint32_t q_0, const uint32_t q_1, const uint32_t q_2, half2 (&dq)[16], int stride, const uint32_t zero) { - const uint32_t c0 = 0x64006400; - const half y8_ = __float2half_rn(1.0f / 8.0f); - const half y64_ = __float2half_rn(1.0f / 64.0f); - const half2 y8 = __halves2half2(y8_, y8_); - const half2 y64 = __halves2half2(y64_, y64_); - const half_uint16 z1_(0xe400 | zero); // half(-1024.0f - zero); - const half z8_ = __hsub(__int2half_rn(-128), __int2half_rn(zero)); - const half z64_ = __hsub(__int2half_rn(-16), __int2half_rn(zero)); - const half2 z1 = __halves2half2(z1_.as_half, z1_.as_half); - const half2 z8 = __halves2half2(z8_, z8_); - const half2 z64 = __halves2half2(z64_, z64_); - - uint32_t qa = q_0; - uint32_t qb = q_1; - uint32_t qc = q_2; - - half2_uint32 q0((qa & 0x00070007) | c0); // half2(q[ 0], q[ 1]) + 1024 - half2_uint32 q1((qa & 0x00380038) | c0); // half2(q[ 2], q[ 3]) * 8 + 1024 - qa >>= 6; - half2_uint32 q2((qa & 0x00070007) | c0); // half2(q[ 4], q[ 5]) + 1024 - half2_uint32 q3((qa & 0x00380038) | c0); // half2(q[ 6], q[ 7]) * 8 + 1024 - half2_uint32 q4((qa & 0x01c001c0) | c0); // half2(q[ 8], q[ 9]) * 64 + 1024 - qa >>= 9; - qa &= 0x00010001; - half2_uint32 q5((qb & 0x00070007) | c0); // half2(q[10], q[11]) + 1024 - half2_uint32 q6((qb & 0x00380038) | c0); // half2(q[12], q[13]) * 8 + 1024 - qb >>= 6; - half2_uint32 q7((qb & 0x00070007) | c0); // half2(q[14], q[15]) + 1024 - half2_uint32 q8((qb & 0x00380038) | c0); // half2(q[16], q[17]) * 8 + 1024 - half2_uint32 q9((qb & 0x01c001c0) | c0); // half2(q[18], q[19]) * 64 + 1024 - qb >>= 8; - qb &= 0x00020002; - half2_uint32 q10((qc & 0x00070007) | c0); // half2(q[20], q[21]) + 1024 - half2_uint32 q11((qc & 0x00380038) | c0); // half2(q[22], q[23]) * 8 + 1024 - qc >>= 6; - half2_uint32 q12((qc & 0x00070007) | c0); // half2(q[24], q[25]) + 1024 - half2_uint32 q13((qc & 0x00380038) | c0); // half2(q[26], q[27]) * 8 + 1024 - half2_uint32 q14((qc & 0x01c001c0) | c0); // half2(q[28], q[29]) * 64 + 1024 - qc >>= 7; - qc &= 0x00040004; - half2_uint32 q15((qa | qb | qc) | c0); - - dq[0] = __hadd2(q0.as_half2, z1); - dq[1] = __hfma2(q1.as_half2, y8, z8); - dq[2] = __hadd2(q2.as_half2, z1); - dq[3] = __hfma2(q3.as_half2, y8, z8); - dq[4] = __hfma2(q4.as_half2, y64, z64); - dq[5] = __hadd2(q5.as_half2, z1); - dq[6] = __hfma2(q6.as_half2, y8, z8); - dq[7] = __hadd2(q7.as_half2, z1); - dq[8] = __hfma2(q8.as_half2, y8, z8); - dq[9] = __hfma2(q9.as_half2, y64, z64); - dq[10] = __hadd2(q10.as_half2, z1); - dq[11] = __hfma2(q11.as_half2, y8, z8); - dq[12] = __hadd2(q12.as_half2, z1); - dq[13] = __hfma2(q13.as_half2, y8, z8); - dq[14] = __hfma2(q14.as_half2, y64, z64); - dq[15] = __hadd2(q15.as_half2, z1); -} - -} // namespace gptq -} // namespace sglang - -#endif diff --git a/python/sglang/kernels/aot/csrc/gemm/gptq/qdq_4.cuh b/python/sglang/kernels/aot/csrc/gemm/gptq/qdq_4.cuh deleted file mode 100644 index c96af4718..000000000 --- a/python/sglang/kernels/aot/csrc/gemm/gptq/qdq_4.cuh +++ /dev/null @@ -1,114 +0,0 @@ -/* -Copied from https://github.com/turboderp/exllamav2 -*/ - -#ifndef _qdq_4_cuh -#define _qdq_4_cuh - -#include "qdq_util.cuh" - -namespace sglang { -namespace gptq { -// Permutation: -// -// 77775555 33331111 66664444 22220000 - -__forceinline__ __device__ void shuffle_4bit_8(uint32_t* q, int stride) { - uint32_t qa = q[0]; - uint32_t qb = 0; - -#pragma unroll - for (int i = 0; i < 4; i++) { - uint32_t qa0 = qa & 0x0f; - uint32_t qa1 = (qa & 0xf0) >> 4; - qa >>= 8; - qb |= (qa1 << (i * 4 + 16)); - qb |= (qa0 << (i * 4)); - } - q[0] = qb; -} - -__forceinline__ __device__ void dequant_4bit_8(const uint32_t q_0, half2 (&dq)[4], int stride, const uint32_t zero) { - const uint32_t c0 = 0x64006400; - const half y16_ = __float2half_rn(1.0f / 16.0f); - const half2 y16 = __halves2half2(y16_, y16_); - const half_uint16 z1_(0xe400 | zero); // half(-1024.0f - zero); - const half z16_ = __hsub(__int2half_rn(-64), __int2half_rn(zero)); - const half2 z1 = __half2half2(z1_.as_half); - const half2 z16 = __half2half2(z16_); - - uint32_t qa = q_0; - half2_uint32 q0((qa & 0x000f000f) | c0); // half2(q[ 0], q[ 1]) + 1024 - half2_uint32 q1((qa & 0x00f000f0) | c0); // half2(q[ 2], q[ 3]) * 16 + 1024 - qa >>= 8; - half2_uint32 q2((qa & 0x000f000f) | c0); // half2(q[ 4], q[ 5]) + 1024 - half2_uint32 q3((qa & 0x00f000f0) | c0); // half2(q[ 6], q[ 7]) * 16 + 1024 - - dq[0] = __hadd2(q0.as_half2, z1); - dq[1] = __hfma2(q1.as_half2, y16, z16); - dq[2] = __hadd2(q2.as_half2, z1); - dq[3] = __hfma2(q3.as_half2, y16, z16); -} - -__forceinline__ __device__ void -dequant_4bit_8_prep_zero_scale(const uint32_t zero, const half scale, half2 (&z1z16)[2], half2 (&y1y16)[2]) { - half_uint16 z1(0xe400 | zero); // half(-1024.0f - zero); - half z16 = __hsub(__int2half_rn(-64), __int2half_rn(zero)); - - half2 scale2 = __half2half2(scale); - - z1z16[0] = __hmul2(scale2, __half2half2(z1.as_half)); - z1z16[1] = __hmul2(scale2, __half2half2(z16)); - - const half y1 = __float2half_rn(1.0f); - const half y16 = __float2half_rn(1.0f / 16.0f); - - y1y16[0] = __hmul2(scale2, __half2half2(y1)); - y1y16[1] = __hmul2(scale2, __half2half2(y16)); -} - -__forceinline__ __device__ void dequant_4bit_8_prep_zero(const uint32_t zero, half2 (&z1z16)[2], half2 (&y1y16)[2]) { - half_uint16 z1(0xe400 | zero); // half(-1024.0f - zero); - half z16 = __hsub(__int2half_rn(-64), __int2half_rn(zero)); - - z1z16[0] = __half2half2(z1.as_half); - z1z16[1] = __half2half2(z16); - - const half y1 = __float2half_rn(1.0f); - const half y16 = __float2half_rn(1.0f / 16.0f); - - y1y16[0] = __half2half2(y1); - y1y16[1] = __half2half2(y16); -} - -__forceinline__ __device__ void -dequant_4bit_8_gptq(const uint32_t q_0, half2 (&dq)[4], half2 (&z1z16)[2], half2 (&y1y16)[2], int stride, bool scaled) { - const uint32_t c0 = 0x64006400; - - uint32_t qa = q_0; - half2_uint32 q0((qa & 0x000f000f) | c0); // half2( q[0] + 1024, q[1] + 1024 ) - half2_uint32 q1((qa & 0x00f000f0) | c0); // half2( q[2] * 16 + 1024, q[3] * 16 + 1024 ) - qa >>= 8; - half2_uint32 q2((qa & 0x000f000f) | c0); // half2( q[4] + 1024, q[5] + 1024 ) - half2_uint32 q3((qa & 0x00f000f0) | c0); // half2( q[6] * 16 + 1024, q[7] * 16 + 1024 ) - - if (scaled) { - dq[0] = __hfma2(q0.as_half2, y1y16[0], - z1z16[0]); // half2( q[0] * s - z * s, q[1] * s - z * s) - dq[1] = __hfma2(q1.as_half2, y1y16[1], - z1z16[1]); // half2( q[2] * s - z * s, q[3] * s - z * s) - dq[2] = __hfma2(q2.as_half2, y1y16[0], z1z16[0]); - dq[3] = __hfma2(q3.as_half2, y1y16[1], z1z16[1]); - } else { - dq[0] = __hadd2(q0.as_half2, z1z16[0]); // half2( q[0] - z, q[1] - z ) - dq[1] = __hfma2(q1.as_half2, y1y16[1], - z1z16[1]); // half2( q[2] - z, q[3] - z ) - dq[2] = __hadd2(q2.as_half2, z1z16[0]); // half2( q[4] - z, q[5] - z ) - dq[3] = __hfma2(q3.as_half2, y1y16[1], - z1z16[1]); // half2( q[6] - z, q[7] - z ) - } -} -} // namespace gptq -} // namespace sglang - -#endif diff --git a/python/sglang/kernels/aot/csrc/gemm/gptq/qdq_8.cuh b/python/sglang/kernels/aot/csrc/gemm/gptq/qdq_8.cuh deleted file mode 100644 index c6a49d6dc..000000000 --- a/python/sglang/kernels/aot/csrc/gemm/gptq/qdq_8.cuh +++ /dev/null @@ -1,30 +0,0 @@ -/* -Copied from https://github.com/turboderp/exllamav2 -*/ - -#ifndef _qdq_8_cuh -#define _qdq_8_cuh - -#include "qdq_util.cuh" - -namespace sglang { -namespace gptq { - -__forceinline__ __device__ void shuffle_8bit_4(uint32_t* q, int stride) {} - -__forceinline__ __device__ void -dequant_8bit_8(const uint32_t q_0, const uint32_t q_1, half2 (&dq)[4], int stride, const uint32_t zero) { - half dqh[8]; - for (int i = 0; i < 4; i++) - dqh[i] = dq_ns(exb(q_0, i * 8, 0xff), zero); - for (int i = 0; i < 4; i++) - dqh[i + 4] = dq_ns(exb(q_1, i * 8, 0xff), zero); - - for (int i = 0; i < 4; i++) - dq[i] = __halves2half2(dqh[i * 2], dqh[i * 2 + 1]); -} - -} // namespace gptq -} // namespace sglang - -#endif diff --git a/python/sglang/kernels/aot/csrc/gemm/gptq/qdq_util.cuh b/python/sglang/kernels/aot/csrc/gemm/gptq/qdq_util.cuh deleted file mode 100644 index 0977269c3..000000000 --- a/python/sglang/kernels/aot/csrc/gemm/gptq/qdq_util.cuh +++ /dev/null @@ -1,53 +0,0 @@ -/* -Copied from https://github.com/turboderp/exllamav2 -*/ - -#ifndef _qdq_util_cuh -#define _qdq_util_cuh - -namespace sglang { -namespace gptq { - -union half2_uint32 { - uint32_t as_uint32; - half2 as_half2; - __device__ half2_uint32(uint32_t val) : as_uint32(val) {} - __device__ half2_uint32(half2 val) : as_half2(val) {} -}; - -union half_uint16 { - uint16_t as_uint16; - half as_half; - __device__ half_uint16(uint16_t val) : as_uint16(val) {} - __device__ half_uint16(half val) : as_half(val) {} -}; - -// Max_scale premultiplied by 1/256 - -__forceinline__ __device__ half dq_scale(const int qs, const half max_scale) { - int qs_i = qs + 1; - half qs_h = __int2half_rn(qs_i * qs_i); - qs_h = __hmul(qs_h, max_scale); - return qs_h; -} - -__forceinline__ __device__ half dq(const int q, const int qzero, const half scale) { - return __hmul(__int2half_rn(q - qzero), scale); -} - -__forceinline__ __device__ half dq_ns(const int q, const int qzero) { - // return __hsub(__int2half_rn(q), __int2half_rn(qzero)); - return __int2half_rn(q - qzero); -} - -__forceinline__ __device__ int exb(const uint32_t q, const int shift, const int mask) { - return (int)((q >> shift) & mask); -} - -__forceinline__ __device__ int exb(const uint32_t q1, const uint32_t q0, const int shift, const int mask) { - return (int)(__funnelshift_rc(q0, q1, shift) & mask); -} - -} // namespace gptq -} // namespace sglang -#endif diff --git a/python/sglang/kernels/aot/include/sgl_kernel_ops.h b/python/sglang/kernels/aot/include/sgl_kernel_ops.h index 4ee916228..846e67efd 100644 --- a/python/sglang/kernels/aot/include/sgl_kernel_ops.h +++ b/python/sglang/kernels/aot/include/sgl_kernel_ops.h @@ -96,21 +96,6 @@ void register_graph_buffers( */ 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( - torch::Tensor const& out, - torch::Tensor const& q_nope, - torch::Tensor const& q_pe, - torch::Tensor const& kv_c_and_k_pe_cache, - torch::Tensor const& seq_lens, - torch::Tensor const& page_table, - torch::Tensor const& workspace, - double sm_scale, - int64_t num_kv_splits = 1 /* Set to 1 to avoid cuda_graph issue by default. */); -int64_t cutlass_mla_get_workspace_size( - int64_t max_seq_len, - int64_t num_batches, - int64_t sm_count = 0, - int64_t num_kv_splits = 1 /* Set to 1 to avoid cuda_graph issue by default. */); /* * From csrc/infllm_v2 @@ -220,7 +205,6 @@ void dsv4_fused_q_indexer_rope_hadamard_quant( /* * From csrc/gemm */ -torch::Tensor awq_dequantize(torch::Tensor qweight, torch::Tensor scales, torch::Tensor qzeros); torch::Tensor int8_scaled_mm( const torch::Tensor& mat_a, const torch::Tensor& mat_b, @@ -257,17 +241,6 @@ void sgl_per_token_group_quant_8bit_v2( const std::optional& masked_m); void sgl_per_token_quant_fp8(at::Tensor input, at::Tensor output_q, at::Tensor output_s); -torch::Tensor gptq_gemm( - torch::Tensor a, - torch::Tensor b_q_weight, - torch::Tensor b_gptq_qzeros, - torch::Tensor b_gptq_scales, - torch::Tensor b_g_idx, - bool use_shuffle, - int64_t bit); - -void gptq_shuffle(torch::Tensor q_weight, torch::Tensor q_perm, int64_t bit); - /* * From csrc/moe */ @@ -643,82 +616,6 @@ void top_k_renorm_probs( void top_p_renorm_probs( at::Tensor probs, at::Tensor renorm_probs, std::optional maybe_top_p_arr, double top_p_val); -namespace flash { -/* - * From fa2 sparse - */ -std::vector mha_fwd_sparse( - at::Tensor& q, // batch_size x seqlen_q x num_heads x head_size - const at::Tensor& k, // batch_size x seqlen_k x num_heads_k x head_size - const at::Tensor& v, // batch_size x seqlen_k x num_heads_k x head_size - const at::Tensor& block_count, - const at::Tensor& block_offset, - const at::Tensor& column_count, - const at::Tensor& column_index, - const std::optional& out_, // batch_size x seqlen_q x num_heads x head_size - const std::optional& alibi_slopes_, // num_heads or batch_size x num_heads - const double p_dropout, - const double softmax_scale, - bool is_causal, - const double softcap, - const bool return_softmax, - std::optional gen_); - -std::vector mha_varlen_fwd_sparse( - at::Tensor& q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i - const at::Tensor& k, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i. - const at::Tensor& v, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i. - const at::Tensor& block_count, - const at::Tensor& block_offset, - const at::Tensor& column_count, - const at::Tensor& column_index, - const c10::optional& out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i - const at::Tensor& cu_seqlens_q, // b+1 - const at::Tensor& cu_seqlens_k, // b+1 - const c10::optional& - seqused_k, // b. If given, only this many elements of each batch element's keys are used. - const c10::optional& alibi_slopes_, // num_heads or b x num_heads - int64_t max_seqlen_q, - const int64_t max_seqlen_k, - const double p_dropout, - const double softmax_scale, - const bool zero_tensors, - bool is_causal, - const double softcap, - const bool return_softmax, - c10::optional gen_); -} // namespace flash - -void convert_vertical_slash_indexes( - torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS] - torch::Tensor& block_offset, // [BATCH, N_HEADS, NUM_ROWS, NNZ_S] - torch::Tensor& column_count, // [BATCH, N_HEADS, NUM_ROWS] - torch::Tensor& column_index, // [BATCH, N_HEADS, NUM_ROWS, NNZ_V] - torch::Tensor q_seqlens, // [BATCH, ] - torch::Tensor kv_seqlens, // [BATCH, ] - torch::Tensor vertical_indexes, // [BATCH, N_HEADS, NNZ_V] - torch::Tensor slash_indexes, // [BATCH, N_HEADS, NNZ_S] - int64_t context_size, - int64_t block_size_M, - int64_t block_size_N, - bool causal); - -void convert_vertical_slash_indexes_mergehead( - torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS] - torch::Tensor& block_offset, // [BATCH, N_HEADS, NUM_ROWS, NNZ_S] - torch::Tensor& column_count, // [BATCH, N_HEADS, NUM_ROWS] - torch::Tensor& column_index, // [BATCH, N_HEADS, NUM_ROWS, NNZ_V] - torch::Tensor q_seqlens, // [BATCH, ] - torch::Tensor kv_seqlens, // [BATCH, ] - torch::Tensor vertical_indexes, // [BATCH, N_HEADS, NNZ_V] - torch::Tensor slash_indexes, // [BATCH, N_HEADS, NNZ_S] - torch::Tensor vertical_indices_count, // [N_HEADS, ] - torch::Tensor slash_indices_count, - int64_t context_size, - int64_t block_size_M, - int64_t block_size_N, - bool causal); - /* * From csrc/grammar */ diff --git a/python/sglang/kernels/aot/python/sgl_kernel/__init__.py b/python/sglang/kernels/aot/python/sgl_kernel/__init__.py index f725f62c2..ef6840eaf 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/__init__.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/__init__.py @@ -23,11 +23,7 @@ else: _preload_cuda_library() from sgl_kernel.allreduce import * - from sgl_kernel.attention import ( - cutlass_mla_decode, - cutlass_mla_get_workspace_size, - merge_state_v2, - ) + from sgl_kernel.attention import merge_state_v2 from sgl_kernel.cutlass_moe import ( cutlass_w4a8_moe_mm, get_cutlass_w4a8_moe_mm_data, @@ -54,10 +50,7 @@ else: es_sm100_mxfp8_blockscaled_grouped_quant, ) from sgl_kernel.gemm import ( - awq_dequantize, fp8_scaled_mm, - gptq_gemm, - gptq_shuffle, int8_scaled_mm, sgl_per_token_group_quant_8bit, sgl_per_token_group_quant_fp8, @@ -150,15 +143,12 @@ else: _DEBUG_EXPORT_NAMES = [ "apply_shuffle_mul_sum", "apply_token_bitmask_inplace_cuda", - "awq_dequantize", "build_tree_kernel_efficient", "causal_conv1d_fwd", "causal_conv1d_update", "concat_mla_absorb_q", "concat_mla_k", "copy_to_gpu_no_ce", - "cutlass_mla_decode", - "cutlass_mla_get_workspace_size", "dsv4_fused_k_norm_rope_flashmla", "dsv4_fused_q_indexer_rope_hadamard_quant", "dsv4_fused_q_norm_rope", @@ -177,8 +167,6 @@ else: "gelu_tanh_and_mul", "gemma_fused_add_rmsnorm", "gemma_rmsnorm", - "gptq_gemm", - "gptq_shuffle", "int8_scaled_mm", "merge_state_v2", "moe_align_block_size", diff --git a/python/sglang/kernels/aot/python/sgl_kernel/attention.py b/python/sglang/kernels/aot/python/sgl_kernel/attention.py index 3351d6d6f..7733b6e64 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/attention.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/attention.py @@ -24,90 +24,3 @@ def merge_state_v2( s_merged = torch.empty_like(s_a) torch.ops.sgl_kernel.merge_state_v2.default(v_a, s_a, v_b, s_b, v_merged, s_merged) return v_merged, s_merged - - -def cutlass_mla_decode( - q_nope: torch.Tensor, - q_pe: torch.Tensor, - kv_c_and_k_pe_cache: torch.Tensor, - seq_lens: torch.Tensor, - page_table: torch.Tensor, - workspace: torch.Tensor, - sm_scale: float, - num_kv_splits: int = 1, # Set to 1 to avoid cuda_graph issue by default. -) -> torch.Tensor: - assert q_nope.ndim == 3, f"q_nope must be a 3D tensor, but got {q_nope.ndim}" - assert q_pe.ndim == 3, f"q_pe must be a 3D tensor, but got {q_pe.ndim}" - assert kv_c_and_k_pe_cache.ndim == 3, ( - f"kv_c_and_k_pe_cache must be a 3D tensor, but got {kv_c_and_k_pe_cache.ndim}" - ) - - B_q, H, D_q_nope = q_nope.shape - B_q_2, H_2, D_q_pe = q_pe.shape - assert (B_q == B_q_2) and (H == H_2) - - _, PAGE_SIZE, D_ckv = kv_c_and_k_pe_cache.shape - - D_latent = 512 - D_rope = 64 - assert D_q_nope == D_latent - assert D_q_pe == D_rope - assert D_ckv == D_latent + D_rope - - MAX_HEADS = 128 - assert H <= MAX_HEADS, f"H must be <= {MAX_HEADS}, but got {H}" - if H < MAX_HEADS: - q_nope_padded = q_nope.new_empty((B_q, MAX_HEADS, D_q_nope)) - q_nope_padded[:, :H] = q_nope - q_nope = q_nope_padded - - q_pe_padded = q_pe.new_empty((B_q, MAX_HEADS, D_q_pe)) - q_pe_padded[:, :H] = q_pe - q_pe = q_pe_padded - - assert len(page_table.shape) == 2 - B_block_table, block_num = page_table.shape - assert B_block_table == B_q - assert block_num > 0, f"block num must be greater than 0, got {block_num}" - assert block_num % (128 / PAGE_SIZE) == 0 - - # TODO(kaixih@nvidia): support fp8 - assert q_nope.dtype in ( - torch.float16, - torch.bfloat16, - ), f"q_nope.dtype needs to be fp16 or bf16 but got {q_nope.dtype}." - assert q_nope.dtype == q_pe.dtype == kv_c_and_k_pe_cache.dtype - assert seq_lens.dtype == torch.int32, ( - f"seq_lens.dtype needs to be int32 but got {seq_lens.dtype}." - ) - assert page_table.dtype == torch.int32, ( - f"page_table.dtype needs to be int32 but got {page_table.dtype}." - ) - - out = q_nope.new_empty((B_q, MAX_HEADS, D_latent)) - - torch.ops.sgl_kernel.cutlass_mla_decode.default( - out, - q_nope, - q_pe, - kv_c_and_k_pe_cache, - seq_lens, - page_table, - workspace, - sm_scale, - num_kv_splits, - ) - return out[:, :H].contiguous() - - -def cutlass_mla_get_workspace_size( - max_seq_len: int, - num_batches: int, - sm_count: int = 0, - num_kv_splits: int = 1, # Set to 1 to avoid cuda_graph issue by default. -) -> int: - assert max_seq_len > 0, f"max_seq_len must be greater than 0, got {max_seq_len}" - assert num_batches > 0, f"num_batches must be greater than 0, got {num_batches}" - return torch.ops.sgl_kernel.cutlass_mla_get_workspace_size.default( - max_seq_len, num_batches, sm_count, num_kv_splits - ) diff --git a/python/sglang/kernels/aot/python/sgl_kernel/gemm.py b/python/sglang/kernels/aot/python/sgl_kernel/gemm.py index fa896625d..6831f3d95 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/gemm.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/gemm.py @@ -3,12 +3,6 @@ from typing import Optional import torch -def awq_dequantize( - qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor -) -> torch.ByteTensor: - return torch.ops.sgl_kernel.awq_dequantize.default(qweight, scales, qzeros) - - def int8_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None): return torch.ops.sgl_kernel.int8_scaled_mm.default( mat_a, @@ -90,22 +84,3 @@ def shuffle_rows(input_tensor, dst2src_map, output_tensor_shape): ) torch.ops.sgl_kernel.shuffle_rows.default(input_tensor, dst2src_map, output_tensor) return output_tensor - - -# GPTQ kernels -def gptq_gemm( - a: torch.Tensor, - b_q_weight: torch.Tensor, - b_gptq_qzeros: torch.Tensor, - b_gptq_scales: torch.Tensor, - b_g_idx: torch.Tensor, - use_shuffle: bool, - bit: int, -) -> torch.Tensor: - return torch.ops.sgl_kernel.gptq_gemm( - a, b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, use_shuffle, bit - ) - - -def gptq_shuffle(q_weight: torch.Tensor, q_perm: torch.Tensor, bit: int) -> None: - torch.torch.ops.sgl_kernel.gptq_shuffle(q_weight, q_perm, bit) diff --git a/python/sglang/kernels/aot/python/sgl_kernel/sparse_flash_attn.py b/python/sglang/kernels/aot/python/sgl_kernel/sparse_flash_attn.py deleted file mode 100644 index 29b2f0405..000000000 --- a/python/sglang/kernels/aot/python/sgl_kernel/sparse_flash_attn.py +++ /dev/null @@ -1,293 +0,0 @@ -from typing import List, Optional, Tuple, Union - -import torch -import torch.nn as nn - - -def maybe_contiguous(x): - return x.contiguous() if x is not None and x.stride(-1) != 1 else x - - -# Sparse attention utils -def convert_vertical_slash_indexes( - q_seqlens: torch.Tensor, # [BATCH, ] - kv_seqlens: torch.Tensor, # [BATCH, ] - vertical_indexes: torch.Tensor, # [BATCH, N_HEADS, NNZ_V] - slash_indexes: torch.Tensor, # [BATCH, N_HEADS, NNZ_S] - context_size: int, - block_size_M: int, - block_size_N: int, - causal: bool = True, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - batch_size = slash_indexes.size(0) - num_heads = slash_indexes.size(1) - nnz_slash = slash_indexes.size(2) - nnz_vertical = vertical_indexes.size(2) - num_rows = (context_size + block_size_M - 1) // block_size_M - - block_count = torch.zeros( - batch_size, num_heads, num_rows, dtype=q_seqlens.dtype, device=q_seqlens.device - ) - block_offset = torch.zeros( - batch_size, - num_heads, - num_rows, - nnz_slash, - dtype=q_seqlens.dtype, - device=q_seqlens.device, - ) - column_count = torch.zeros( - batch_size, num_heads, num_rows, dtype=q_seqlens.dtype, device=q_seqlens.device - ) - column_index = torch.zeros( - batch_size, - num_heads, - num_rows, - nnz_vertical, - dtype=q_seqlens.dtype, - device=q_seqlens.device, - ) - - torch.ops.sgl_kernel.convert_vertical_slash_indexes.default( - block_count, - block_offset, - column_count, - column_index, - q_seqlens, - kv_seqlens, - vertical_indexes, - slash_indexes, - context_size, - block_size_M, - block_size_N, - causal, - ) - return block_count, block_offset, column_count, column_index - - -def convert_vertical_slash_indexes_mergehead( - q_seqlens: torch.Tensor, # [BATCH, ] - kv_seqlens: torch.Tensor, # [BATCH, ] - vertical_indexes: torch.Tensor, # [BATCH, N_HEADS, NNZ_V] - slash_indexes: torch.Tensor, # [BATCH, N_HEADS, NNZ_S] - # [N_HEADS] : different head use different number of indices - vertical_indices_count: torch.Tensor, - slash_indices_count: torch.Tensor, - context_size: int, - block_size_M: int, - block_size_N: int, - causal: bool = True, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - batch_size = slash_indexes.size(0) - num_heads = slash_indexes.size(1) - nnz_slash = slash_indexes.size(2) - nnz_vertical = vertical_indexes.size(2) - num_rows = (context_size + block_size_M - 1) // block_size_M - - block_count = torch.empty( - batch_size, num_heads, num_rows, dtype=q_seqlens.dtype, device=q_seqlens.device - ) - block_offset = torch.empty( - batch_size, - num_heads, - num_rows, - nnz_slash, - dtype=q_seqlens.dtype, - device=q_seqlens.device, - ) - column_count = torch.empty( - batch_size, num_heads, num_rows, dtype=q_seqlens.dtype, device=q_seqlens.device - ) - column_index = torch.empty( - batch_size, - num_heads, - num_rows, - nnz_vertical, - dtype=q_seqlens.dtype, - device=q_seqlens.device, - ) - - torch.ops.sgl_kernel.convert_vertical_slash_indexes_mergehead.default( - block_count, - block_offset, - column_count, - column_index, - q_seqlens, - kv_seqlens, - vertical_indexes, - slash_indexes, - vertical_indices_count, - slash_indices_count, - context_size, - block_size_M, - block_size_N, - causal, - ) - return block_count, block_offset, column_count, column_index - - -def sparse_attn_func( - q, - k, - v, - block_count, - block_offset, - column_count, - column_index, - dropout_p=0.0, - softmax_scale=None, - causal=False, - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - *, - return_softmax_lse=False, - out=None, -): - """Compute attention with vertical and slash sparsity patterns. - Most Arguments are the same with the flash_attn_func interface, except for 4 extra args: - block_count and block_offset for slash sparsity patterns, and - column_count and column_index for vertical sparsity patterns. - For more details please refer to Appendix C.4.2 of paper https://arxiv.org/abs/2407.02490. - - Arguments: - q: (batch_size, seqlen, nheads, headdim) - k: (batch_size, seqlen, nheads_k, headdim) - v: (batch_size, seqlen, nheads_k, headdim) - block_count: (batch_size, nheads, cdiv(seqlen, BLOCK_M)) - block_offset: (batch_size, nheads, cdiv(seqlen, BLOCK_M), NNZ_S) - column_count: (batch_size, nheads, cdiv(seqlen, BLOCK_M)) - column_index: (batch_size, nheads, cdiv(seqlen, BLOCK_M), NNZ_V) - dropout_p: float. Dropout probability. - softmax_scale: float. The scaling of QK^T before applying softmax. - Default to 1 / sqrt(headdim). - causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). - alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of - (-alibi_slope * |i + seqlen_k - seqlen_q - j|) - is added to the attention score of query i and key j. - deterministic: bool. Whether to use the deterministic implementation of the backward pass, - which is slightly slower and uses more memory. The forward pass is always deterministic. - return_attn_probs: bool. Whether to return the attention probabilities. This option is for - testing only. The returned probabilities are not guaranteed to be correct - (they might not have the right scaling). - Return: - out: (batch_size, seqlen, nheads, headdim). - softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The - logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax - normalization factor). - """ - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse = torch.ops.sgl_kernel.fwd_sparse.default( - q, - k, - v, - block_count, - block_offset, - column_count, - column_index, - out, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - softcap, - return_attn_probs and dropout_p > 0, - None, - ) - return (out, softmax_lse) if return_softmax_lse else out - - -def sparse_attn_varlen_func( - q, - k, - v, - block_count, - block_offset, - column_count, - column_index, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - dropout_p=0.0, - softmax_scale=None, - causal=False, - softcap=0.0, # 0.0 means deactivated - alibi_slopes=None, - deterministic=False, - return_attn_probs=False, - *, - return_softmax_lse=False, - out=None, -): - """Compute attention with vertical and slash sparsity patterns. - Most Arguments are the same with the flash_attn_varlen_func interface, except for 4 extra args: - block_count and block_offset for slash sparsity patterns, and - column_count and column_index for vertical sparsity patterns. - For more details please refer to Appendix C.4.2 of paper https://arxiv.org/abs/2407.02490. - - Arguments: - q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. - k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. - block_count: (batch_size, nheads, cdiv(seqlen, BLOCK_M)) - block_offset: (batch_size, nheads, cdiv(seqlen, BLOCK_M), NNZ_S) - column_count: (batch_size, nheads, cdiv(seqlen, BLOCK_M)) - column_index: (batch_size, nheads, cdiv(seqlen, BLOCK_M), NNZ_V) - cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into q. - cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths - of the sequences in the batch, used to index into kv. - max_seqlen_q: int. Maximum query sequence length in the batch. - max_seqlen_k: int. Maximum key sequence length in the batch. - dropout_p: float. Dropout probability. - softmax_scale: float. The scaling of QK^T before applying softmax. - Default to 1 / sqrt(headdim). - causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). - softcap: float. Anything > 0 activates softcapping attention. - alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of - (-alibi_slope * |i + seqlen_k - seqlen_q - j|) - is added to the attention score of query i and key j. - deterministic: bool. Whether to use the deterministic implementation of the backward pass, - which is slightly slower and uses more memory. The forward pass is always deterministic. - return_attn_probs: bool. Whether to return the attention probabilities. This option is for - testing only. The returned probabilities are not guaranteed to be correct - (they might not have the right scaling). - Return: - out: (total, nheads, headdim). - softmax_lse [optional, if return_softmax_lse=True]: (nheads, total_q_seqlen). The - logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax - normalization factor). - """ - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - - q, k, v = [maybe_contiguous(x) for x in (q, k, v)] - out, softmax_lse = torch.ops.sgl_kernel.varlen_fwd_sparse.default( - q, - k, - v, - block_count, - block_offset, - column_count, - column_index, - out, - cu_seqlens_q, - cu_seqlens_k, - None, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - False, - causal, - softcap, - return_attn_probs and dropout_p > 0, - None, - ) - return (out, softmax_lse) if return_softmax_lse else out diff --git a/python/sglang/kernels/aot/setup_musa.py b/python/sglang/kernels/aot/setup_musa.py index 504d6be22..8bea76171 100644 --- a/python/sglang/kernels/aot/setup_musa.py +++ b/python/sglang/kernels/aot/setup_musa.py @@ -94,7 +94,6 @@ sources = [ "csrc/speculative/packbit.cu", "csrc/speculative/speculative_sampling.cu", "csrc/kvcacheio/transfer.cu", - "csrc/gemm/awq_kernel.cu", "csrc/gemm/per_token_quant_fp8.cu", "csrc/gemm/per_token_group_quant_8bit.cu", "csrc/gemm/per_token_group_quant_8bit_v2.cu", diff --git a/python/sglang/kernels/aot/tests/test_awq_dequant.py b/python/sglang/kernels/aot/tests/test_awq_dequant.py deleted file mode 100644 index ce95f5a72..000000000 --- a/python/sglang/kernels/aot/tests/test_awq_dequant.py +++ /dev/null @@ -1,116 +0,0 @@ -import itertools -import sys -from typing import Optional, Tuple - -import pytest -import torch -from sgl_kernel import awq_dequantize - - -def reverse_awq_order(t: torch.Tensor): - bits = 4 - AWQ_REVERSE_ORDER = [0, 4, 1, 5, 2, 6, 3, 7] - reverse_order_tensor = torch.arange( - t.shape[-1], - dtype=torch.int32, - device=t.device, - ) - reverse_order_tensor = reverse_order_tensor.view(-1, 32 // bits) - reverse_order_tensor = reverse_order_tensor[:, AWQ_REVERSE_ORDER] - reverse_order_tensor = reverse_order_tensor.view(-1) - - t = t[:, reverse_order_tensor] & 0xF - return t - - -# qweights - [R , C // 8], int32 -# scales - [R // G, C ], float16 -# zeros - [R // G, C // 8], int32 -def awq_dequantize_torch( - qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor, group_size: int -) -> torch.Tensor: - - if group_size == -1: - group_size = qweight.shape[0] - - bits = 4 - shifts = torch.arange(0, 32, bits, device=qzeros.device) - - iweights = torch.bitwise_right_shift(qweight[:, :, None], shifts[None, None, :]).to( - torch.int8 - ) - - iweights = iweights.view(iweights.shape[0], -1) - - zeros = torch.bitwise_right_shift(qzeros[:, :, None], shifts[None, None, :]).to( - torch.int8 - ) - zeros = zeros.view(qzeros.shape[0], -1) - zeros = reverse_awq_order(zeros) - - iweights = reverse_awq_order(iweights) - - iweights = torch.bitwise_and(iweights, (2**bits) - 1) - zeros = torch.bitwise_and(zeros, (2**bits) - 1) - - scales = scales.repeat_interleave(group_size, dim=0) - zeros = zeros.repeat_interleave(group_size, dim=0) - return (iweights - zeros) * scales - - -def sglang_awq_dequantize( - qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor -) -> torch.Tensor: - return awq_dequantize(qweight, scales, qzeros) - - -@pytest.mark.parametrize( - "qweight_row,qweight_col,is_bf16_act", - list( - itertools.product( - [3584, 18944, 128, 256, 512, 1024, 1536], - [448, 576, 4736, 16, 32, 64, 128, 72], - [True, False], - ) - ), -) -def test_awq_dequant_compare_implementations( - qweight_row: int, qweight_col: int, is_bf16_act: bool -): - device = torch.device("cuda") - qweight = torch.randint( - 0, - torch.iinfo(torch.int32).max, - (qweight_row, qweight_col), - dtype=torch.int32, - device=device, - ) - group_size = qweight_row - scales_row = qweight_row // group_size - scales_col = qweight_col * 8 - - if is_bf16_act: - scales = torch.rand(scales_row, scales_col, dtype=torch.bfloat16, device=device) - else: - scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device) - - qzeros = torch.randint( - 0, - torch.iinfo(torch.int32).max, - (scales_row, qweight_col), - dtype=torch.int32, - device=device, - ) - - # Run both implementations - torch_out = awq_dequantize_torch(qweight, scales, qzeros, group_size) - sglang_out = sglang_awq_dequantize(qweight, scales, qzeros) - - # Compare results - torch.testing.assert_close( - torch_out.to(torch.float32), sglang_out.to(torch.float32), rtol=1e-3, atol=1e-5 - ) - - -if __name__ == "__main__": - sys.exit(pytest.main([__file__])) diff --git a/python/sglang/kernels/aot/tests/test_cutlass_mla.py b/python/sglang/kernels/aot/tests/test_cutlass_mla.py deleted file mode 100644 index 6f2ec7b81..000000000 --- a/python/sglang/kernels/aot/tests/test_cutlass_mla.py +++ /dev/null @@ -1,106 +0,0 @@ -import sys - -import pytest -import torch -import torch.nn.functional as F -from sgl_kernel import cutlass_mla_decode, cutlass_mla_get_workspace_size -from torch import Tensor - -# Disable tests on SM103 until the accuracy issues are fixed. -if torch.cuda.get_device_capability() != (10, 0): - pytest.skip( - reason="Cutlass MLA Requires compute capability of 10.", - allow_module_level=True, - ) - - -def ref_mla( - out: Tensor, # (bs, num_heads, v_head_dim) - query: Tensor, # (bs, num_heads, head_dim) - kv_cache: Tensor, # (num_blocks, block_size, head_dim) - scale: float, - block_tables: Tensor, # (bs, max_num_blocks) - seq_lens: Tensor, # (bs,) -): - bs, num_heads, v_head_dim = out.shape - head_dim = query.shape[2] - - for i in range(bs): - # gather and flatten KV-cache - kv = kv_cache[block_tables[i]] # (max_num_blocks, block_size, head_dim) - kv = kv.view(1, -1, head_dim)[:, : seq_lens[i]] # (1, seq_len, head_dim) - v = kv[:, :, :v_head_dim] - - q = query[i].view(num_heads, 1, head_dim) - o = F.scaled_dot_product_attention(q, kv, v, scale=scale, enable_gqa=True) - out[i] = o.view(num_heads, v_head_dim) - - return out - - -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) -@pytest.mark.parametrize("mean_seq_len", [128, 1024, 4096]) -@pytest.mark.parametrize("bs", [1, 2, 4]) -@pytest.mark.parametrize("varlen", [False, True]) -@pytest.mark.parametrize("block_size", [1, 16, 64, 128]) -@pytest.mark.parametrize("num_heads", [16, 32, 64, 128]) -@pytest.mark.parametrize("num_kv_splits", [-1, 1]) -def test_cutlass_mla_decode( - dtype: torch.dtype, - mean_seq_len: int, - bs: int, - varlen: bool, - block_size: int, - num_heads: int, - num_kv_splits: int, -): - torch.set_default_dtype(dtype) - torch.set_default_device("cuda") - torch.manual_seed(42) - - d = 576 - h_q = num_heads - dv = 512 - - q_nope_dim = 128 - q_pe_dim = 64 - scale = (q_nope_dim + q_pe_dim) ** (-0.5) - if varlen: - seq_lens = torch.empty(bs).normal_(mean_seq_len, mean_seq_len / 2) - seq_lens = seq_lens.clip(2).to(torch.int32) - else: - seq_lens = torch.full((bs,), mean_seq_len, dtype=torch.int32) - max_seq_len = seq_lens.max().item() - block_num = (max_seq_len + block_size - 1) // block_size - - # Pad block_num so that small blocks can be packed into full 128-sized CUTLASS tiles. - # One 128-wide tile can hold (128 // block_size) small blocks. - pack_factor = 128 // block_size - block_num = ((block_num + pack_factor - 1) // pack_factor) * pack_factor - - # Lager q values to detect split kv error - q = torch.randn(bs, h_q, d) * 100.0 - block_table = torch.randint(0, bs * block_num, (bs, block_num), dtype=torch.int32) - - kv_cache = torch.randn(block_table.numel(), block_size, d) - - workspace_size = cutlass_mla_get_workspace_size( - block_num * block_size, bs, num_kv_splits=num_kv_splits - ) - workspace = torch.empty(workspace_size, device="cuda", dtype=torch.uint8) - - q_nope = torch.empty((h_q, bs, dv)).transpose(0, 1) - q_nope.copy_(q[:, :, :dv]) - q_pe = q[:, :, dv:].clone() - - out_ref = q.new_zeros(bs, h_q, dv) - ref_mla(out_ref, q, kv_cache, scale, block_table, seq_lens) - out = cutlass_mla_decode( - q_nope, q_pe, kv_cache, seq_lens, block_table, workspace, scale, num_kv_splits - ) - - torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2) - - -if __name__ == "__main__": - sys.exit(pytest.main([__file__])) diff --git a/python/sglang/kernels/aot/tests/test_flash_attn_sparse.py b/python/sglang/kernels/aot/tests/test_flash_attn_sparse.py deleted file mode 100644 index 0932ca18a..000000000 --- a/python/sglang/kernels/aot/tests/test_flash_attn_sparse.py +++ /dev/null @@ -1,491 +0,0 @@ -import math -import sys -from typing import List, Optional - -import pytest -import torch -from einops import rearrange, repeat -from sgl_kernel.sparse_flash_attn import ( - convert_vertical_slash_indexes, - convert_vertical_slash_indexes_mergehead, - sparse_attn_func, -) -from test_flash_attention import construct_local_mask, is_fa3_supported - - -def ref_attn( - q, - k, - v, - query_padding_mask=None, - key_padding_mask=None, - attn_bias=None, - dropout_p=0.0, - dropout_mask=None, - causal=False, - window_size=(-1, -1), # -1 means infinite window size - softcap=0.0, - upcast=True, - reorder_ops=False, - key_leftpad=None, -): - """ - Arguments: - q: (batch_size, seqlen_q, nheads, head_dim) - k: (batch_size, seqlen_k, nheads_k, head_dim) - v: (batch_size, seqlen_k, nheads_k, head_dim) - query_padding_mask: (batch_size, seqlen_q) - key_padding_mask: (batch_size, seqlen_k) - attn_bias: broadcastable to (batch_size, nheads, seqlen_q, seqlen_k) - dropout_p: float - dropout_mask: (batch_size, nheads, seqlen_q, seqlen_k) - causal: whether to apply causal masking - window_size: (int, int), left and right window size - upcast: whether to cast all inputs to fp32, do all computation in fp32, then cast - output back to fp16/bf16. - reorder_ops: whether to change the order of operations (scaling k instead of scaling q, etc.) - without changing the math. This is to estimate the numerical error from operation - reordering. - Output: - output: (batch_size, seqlen_q, nheads, head_dim) - lse: (batch_size, nheads, seqlen_q) - """ - if causal: - window_size = (window_size[0], 0) - dtype_og = q.dtype - if upcast: - q, k, v = q.float(), k.float(), v.float() - seqlen_q, seqlen_k = q.shape[1], k.shape[1] - k = repeat(k, "b s h d -> b s (h g) d", g=q.shape[2] // k.shape[2]) - v = repeat(v, "b s h d -> b s (h g) d", g=q.shape[2] // v.shape[2]) - d = q.shape[-1] - if not reorder_ops: - scores = torch.einsum("bthd,bshd->bhts", q / math.sqrt(d), k) - else: - scores = torch.einsum("bthd,bshd->bhts", q, k / math.sqrt(d)) - - lse_ref = scores.logsumexp(dim=-1) - - if softcap > 0: - scores = scores / softcap - scores = scores.tanh() - scores = scores * softcap - if key_padding_mask is not None: - scores.masked_fill_( - rearrange(~key_padding_mask, "b s -> b 1 1 s"), float("-inf") - ) - if window_size[0] >= 0 or window_size[1] >= 0: - local_mask = construct_local_mask( - seqlen_q, - seqlen_k, - window_size, - query_padding_mask, - key_padding_mask, - q.device, - key_leftpad=key_leftpad, - ) - scores.masked_fill_(local_mask, float("-inf")) - if attn_bias is not None: - scores = scores + attn_bias - attention = torch.softmax(scores, dim=-1).to(v.dtype) - # Some rows might be completely masked out so we fill them with zero instead of NaN - if window_size[0] >= 0 or window_size[1] >= 0: - attention = attention.masked_fill( - torch.all(local_mask, dim=-1, keepdim=True), 0.0 - ) - # We want to mask here so that the attention matrix doesn't have any NaNs - # Otherwise we'll get NaN in dV - if query_padding_mask is not None: - attention = attention.masked_fill( - rearrange(~query_padding_mask, "b s -> b 1 s 1"), 0.0 - ) - dropout_scaling = 1.0 / (1 - dropout_p) - # attention_drop = attention.masked_fill(~dropout_mask, 0.0) * dropout_scaling - # output = torch.einsum('bhts,bshd->bthd', attention_drop , v) - if dropout_mask is not None: - attention_drop = attention.masked_fill(~dropout_mask, 0.0) - else: - attention_drop = attention - output = torch.einsum("bhts,bshd->bthd", attention_drop, v * dropout_scaling) - if query_padding_mask is not None: - output.masked_fill_(rearrange(~query_padding_mask, "b s -> b s 1 1"), 0.0) - - return output.to(dtype=dtype_og), lse_ref - - -def ref_paged_attn( - query: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - query_lens: List[int], - kv_lens: List[int], - block_tables: torch.Tensor, - scale: float, - sliding_window: Optional[int] = None, - soft_cap: Optional[float] = None, -) -> torch.Tensor: - num_seqs = len(query_lens) - block_tables = block_tables.cpu().numpy() - _, block_size, num_kv_heads, head_size = key_cache.shape - - outputs: List[torch.Tensor] = [] - start_idx = 0 - for i in range(num_seqs): - query_len = query_lens[i] - kv_len = kv_lens[i] - # clone to avoid clobbering the query tensor - q = query[start_idx : start_idx + query_len].clone() - q *= scale - - num_kv_blocks = (kv_len + block_size - 1) // block_size - block_indices = block_tables[i, :num_kv_blocks] - - k = key_cache[block_indices].view(-1, num_kv_heads, head_size) - k = k[:kv_len] - v = value_cache[block_indices].view(-1, num_kv_heads, head_size) - v = v[:kv_len] - - if q.shape[1] != k.shape[1]: - k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1) - v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1) - attn = torch.einsum("qhd,khd->hqk", q, k).float() - empty_mask = torch.ones(query_len, kv_len) - mask = torch.triu(empty_mask, diagonal=kv_len - query_len + 1).bool() - if sliding_window is not None: - sliding_window_mask = ( - torch.triu( - empty_mask, diagonal=kv_len - (query_len + sliding_window) + 1 - ) - .bool() - .logical_not() - ) - mask |= sliding_window_mask - if soft_cap is not None: - attn = soft_cap * torch.tanh(attn / soft_cap) - attn.masked_fill_(mask, float("-inf")) - attn = torch.softmax(attn, dim=-1).to(v.dtype) - out = torch.einsum("hqk,khd->qhd", attn, v) - - outputs.append(out) - start_idx += query_len - - return torch.cat(outputs, dim=0) - - -@pytest.mark.skipif( - not is_fa3_supported(), - reason="flash_attn at sgl-kernel is only supported on sm90 or sm80", -) -@pytest.mark.parametrize("batch_size", [1, 2]) -@pytest.mark.parametrize( - "seq_lens", - [ - (1, 1), - (1, 1024), - (1, 2048), - (1023, 2049), - (1023, 1023), - (32, 32), - (65, 65), - (129, 129), - ], -) -@pytest.mark.parametrize("num_heads", [1, 2, 4]) -@pytest.mark.parametrize("head_size", [128]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("NNZ_S", [0, 1, 2, 3, 7, 15, 32]) -@torch.inference_mode() -def test_sparse_attention( - batch_size, - seq_lens, - num_heads, - head_size, - dtype, - NNZ_S, -) -> None: - torch.set_default_device("cuda") - torch.cuda.manual_seed_all(0) - block_size_M = 64 - block_size_N = 64 - seqlen_q, seqlen_k = seq_lens - q = torch.randn( - batch_size, seqlen_q, num_heads, head_size, dtype=dtype, requires_grad=False - ) - k = torch.randn( - batch_size, seqlen_k, num_heads, head_size, dtype=dtype, requires_grad=False - ) - v = torch.randn( - batch_size, seqlen_k, num_heads, head_size, dtype=dtype, requires_grad=False - ) - NUM_ROWS = (seqlen_q + block_size_M - 1) // block_size_M - if NNZ_S * block_size_N > seqlen_k: - return - NNZ_V = seqlen_k - NNZ_S * block_size_N - block_count = torch.tensor( - [NNZ_S] * batch_size * NUM_ROWS * num_heads, dtype=torch.int32 - ).reshape(batch_size, num_heads, NUM_ROWS) - column_count = torch.tensor( - [NNZ_V] * batch_size * NUM_ROWS * num_heads, dtype=torch.int32 - ).reshape(batch_size, num_heads, NUM_ROWS) - block_offset = torch.tensor( - [[i * block_size_N for i in range(NNZ_S)]] * batch_size * NUM_ROWS * num_heads, - dtype=torch.int32, - ).reshape(batch_size, num_heads, NUM_ROWS, NNZ_S) - column_index = torch.tensor( - [[NNZ_S * block_size_N + i for i in range(NNZ_V)]] - * batch_size - * NUM_ROWS - * num_heads, - dtype=torch.int32, - ).reshape(batch_size, num_heads, NUM_ROWS, NNZ_V) - out, lse = sparse_attn_func( - q, - k, - v, - block_count, - block_offset, - column_count, - column_index, - return_softmax_lse=True, - ) - - ref_out, ref_lse = ref_attn(q, k, v) - - ( - torch.testing.assert_close(out, ref_out, atol=2e-2, rtol=1e-2), - f"{torch.max(torch.abs(out - ref_out))}", - ) - ( - torch.testing.assert_close(lse, ref_lse, atol=2e-2, rtol=1e-2), - f"{torch.max(torch.abs(lse - ref_lse))}", - ) - - -# sparse attention utils -# origin -@pytest.mark.skipif( - not is_fa3_supported(), - reason="flash_attn at sgl-kernel is only supported on sm90 or sm80", -) -@pytest.mark.parametrize("causal", [True, False]) -def test_convert_vertical_slash_indexes(causal): - # Prepare small, hand-checkable inputs - q_seqlens = torch.tensor([4], dtype=torch.int32, device="cuda") # [BATCH] - kv_seqlens = torch.tensor([4], dtype=torch.int32, device="cuda") - vertical_indexes = torch.tensor( - [[[1, 3]]], dtype=torch.int32, device="cuda" - ) # [BATCH, N_HEADS, NNZ_V] - slash_indexes = torch.tensor( - [[[2]]], dtype=torch.int32, device="cuda" - ) # [BATCH, N_HEADS, NNZ_S] - context_size = 4 - block_size_M = 2 - block_size_N = 2 - - # Call your CUDA kernel wrapper - block_count, block_offset, column_count, column_index = ( - convert_vertical_slash_indexes( - q_seqlens, - kv_seqlens, - vertical_indexes, - slash_indexes, - context_size, - block_size_M, - block_size_N, - causal=causal, - ) - ) - - # Manually create expected outputs for this input - # There are 2 rows (blocks): row0 (tokens 0-1), row1 (tokens 2-3) - # Fill these expected tensors based on your CUDA kernel's logic - # For demonstration, we assume: - # - block_count: how many slash indices fall into each block - # - block_offset: the value of those indices - # - column_count: number of valid vertical indices per block - # - column_index: the actual vertical indices - - expected_column_index = torch.tensor( - [[[[0, 0], [0, 0]]]], dtype=torch.int32, device="cuda" - ) - - # If causal=False, update these tensors according to expected behavior - if not causal: - # Update these values if your kernel produces different output in non-causal mode - expected_column_index = torch.tensor( - [[[[1, 0], [1, 3]]]], dtype=torch.int32, device="cuda" - ) - - # Assert that outputs match expectations - assert torch.equal(column_index, expected_column_index) - - -# mergehead -@pytest.mark.skipif( - not is_fa3_supported(), - reason="flash_attn at sgl-kernel is only supported on sm90 or sm80", -) -@pytest.mark.parametrize("causal", [True, False]) -def test_convert_vertical_slash_indexes_mergehead(causal): - # Prepare small, hand-checkable inputs for mergehead version - q_seqlens = torch.tensor([4], dtype=torch.int32, device="cuda") - kv_seqlens = torch.tensor([4], dtype=torch.int32, device="cuda") - vertical_indexes = torch.tensor( - [ - [ - [1, 3], # head 0 - [2, 0], # head 1 - ] - ], - dtype=torch.int32, - device="cuda", - ) # [BATCH, N_HEADS, NNZ_V] - slash_indexes = torch.tensor( - [ - [ - [2, 0], # head 0 - [1, 3], # head 1 - ] - ], - dtype=torch.int32, - device="cuda", - ) # [BATCH, N_HEADS, NNZ_S] - vertical_indices_count = torch.tensor([2, 1], dtype=torch.int32, device="cuda") - slash_indices_count = torch.tensor([1, 2], dtype=torch.int32, device="cuda") - context_size = 4 - block_size_M = 2 - block_size_N = 2 - - # Call your CUDA kernel wrapper - block_count, block_offset, column_count, column_index = ( - convert_vertical_slash_indexes_mergehead( - q_seqlens, - kv_seqlens, - vertical_indexes, - slash_indexes, - vertical_indices_count, - slash_indices_count, - context_size, - block_size_M, - block_size_N, - causal=causal, - ) - ) - - # column_index is torch.empty-backed; only entries before column_count are valid. - expected_column_count = torch.zeros((1, 2, 2), dtype=torch.int32, device="cuda") - expected_column_index = [[[[], []], [[], []]]] - - if not causal: - expected_column_count = torch.tensor( - [[[1, 2], [1, 1]]], dtype=torch.int32, device="cuda" - ) - expected_column_index = [[[[1], [1, 3]], [[2], [2]]]] - - assert torch.equal(column_count, expected_column_count) - for batch_idx, batch_expected in enumerate(expected_column_index): - for head_idx, head_expected in enumerate(batch_expected): - for row_idx, expected_values in enumerate(head_expected): - count = int(column_count[batch_idx, head_idx, row_idx].item()) - actual = column_index[batch_idx, head_idx, row_idx, :count].tolist() - assert actual == expected_values - - -# skip cause use fa2 for test -# @pytest.mark.parametrize("seq_lens", [[(1024, 1328)], -# [(1024, 1328), (1, 2048)], -# [(1025, 1328), (2, 2048)], -# [(1025, 2049), (2, 1281)], -# ]) -# @pytest.mark.parametrize("head_size", [128]) -# @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -# @torch.inference_mode() -# def test_sparse_attention_varlen( -# seq_lens, -# head_size, -# dtype, -# ) -> None: -# torch.set_default_device("cuda") -# torch.cuda.manual_seed_all(0) -# block_size_M = 64 -# block_size_N = 64 -# num_seqs = len(seq_lens) -# query_lens = [x[0] for x in seq_lens] -# kv_lens = [x[1] for x in seq_lens] -# num_heads = 1 -# query = torch.randn(sum(query_lens), -# num_heads, -# head_size, -# dtype=dtype) -# key = torch.randn(sum(kv_lens), -# num_heads, -# head_size, -# dtype=dtype) -# value = torch.randn_like(key) -# cu_query_lens = torch.tensor([0] + query_lens, -# dtype=torch.int32).cumsum(dim=0, -# dtype=torch.int32) -# cu_kv_lens = torch.tensor([0] + kv_lens, -# dtype=torch.int32).cumsum(dim=0, -# dtype=torch.int32) -# max_query_len = max(query_lens) -# max_kv_len = max(kv_lens) - -# NUM_ROWS = (max_query_len + block_size_M - 1) // block_size_M -# NNZ_S = 20 -# NNZ_V = 2048 -# batch_size = len(query_lens) - -# block_counts = [] -# column_counts = [] -# block_offsets = [] -# column_indices = [] -# for b in range(batch_size): -# block_counts.append(torch.tensor([NNZ_S] * NUM_ROWS * num_heads, dtype=torch.int32).reshape(num_heads, NUM_ROWS)) -# columns = kv_lens[b] - NNZ_S * block_size_N -# column_counts.append(torch.tensor([columns] * NUM_ROWS * num_heads, dtype=torch.int32).reshape(num_heads, NUM_ROWS)) -# block_offsets.append(torch.tensor([[i * block_size_N for i in range(NNZ_S)]] * NUM_ROWS * num_heads, dtype=torch.int32).reshape(num_heads, NUM_ROWS, NNZ_S)) -# column_indices.append(torch.tensor([[NNZ_S * block_size_N + i for i in range(NNZ_V)]] * NUM_ROWS * num_heads, dtype=torch.int32).reshape(num_heads, NUM_ROWS, NNZ_V)) -# block_count = torch.concat(block_counts).reshape(batch_size, num_heads, NUM_ROWS) -# column_count = torch.concat(column_counts).reshape(batch_size, num_heads, NUM_ROWS) -# block_offset = torch.concat(block_offsets).reshape(batch_size, num_heads, NUM_ROWS, NNZ_S) -# column_index = torch.concat(column_indices).reshape(batch_size, num_heads, NUM_ROWS, NNZ_V) -# out, lse = sparse_attn_varlen_func( -# query, -# key, -# value, -# block_count, -# block_offset, -# column_count, -# column_index, -# cu_seqlens_q=cu_query_lens, -# cu_seqlens_k=cu_kv_lens, -# max_seqlen_q=max_query_len, -# max_seqlen_k=max_kv_len, -# return_softmax_lse=True, -# ) - -# max_num_blocks_per_seq = (max_kv_len + 2048 - 1) // 2048 -# block_tables = torch.randint(0, -# 2048, -# (len(query_lens), max_num_blocks_per_seq), -# dtype=torch.int32) -# scale = head_size**-0.5 - -# ref_out, ref_lse, _ = ref_paged_attn( -# query, -# key, -# value, -# query_lens=query_lens, -# kv_lens=kv_lens, -# block_tables=block_tables, -# scale=scale -# ) - -# torch.testing.assert_close(out, ref_out, atol=2e-2, rtol=1e-2), \ -# f"{torch.max(torch.abs(out - ref_out))}" -# torch.testing.assert_close(lse, ref_lse, atol=2e-2, rtol=1e-2), \ -# f"{torch.max(torch.abs(lse - ref_lse))}" - -if __name__ == "__main__": - sys.exit(pytest.main([__file__])) diff --git a/python/sglang/kernels/aot/tests/test_gptq_kernel.py b/python/sglang/kernels/aot/tests/test_gptq_kernel.py deleted file mode 100644 index e7596e379..000000000 --- a/python/sglang/kernels/aot/tests/test_gptq_kernel.py +++ /dev/null @@ -1,133 +0,0 @@ -import sys - -import pytest -import torch -from sgl_kernel import gptq_gemm - -from sglang.srt.layers.quantization.utils import pack_cols, pack_rows - - -def torch_dequantize(q_weight, q_zeros, scales, g_idx, use_shuffle, bit, K, N): - assert bit == 4, "Reference dequantization only supports 4-bit" - group_size = K // scales.shape[0] - pack_factor = 32 // bit - - # unpack q_weight: (K//pack_factor, N) -> (K, N) - unpacked_q_weight = torch.empty( - q_weight.shape[0] * pack_factor, - q_weight.shape[1], - dtype=torch.uint8, - device=q_weight.device, - ) - for i in range(pack_factor): - unpacked_q_weight[i::pack_factor, :] = (q_weight >> (i * 4)) & 0x0F - - # unpack q_zeros: (num_groups, N//pack_factor) -> (num_groups, N) - unpacked_q_zeros = torch.empty( - q_zeros.shape[0], - q_zeros.shape[1] * pack_factor, - dtype=torch.uint8, - device=q_zeros.device, - ) - for i in range(pack_factor): - unpacked_q_zeros[:, i::pack_factor] = (q_zeros >> (i * 4)) & 0x0F - - unpacked_q_zeros += 1 - unpacked_q_zeros = unpacked_q_zeros.to(scales.dtype) - - scale_zeros = unpacked_q_zeros * scales # (num_groups, N) - - current_g_idx = torch.tensor( - [i // group_size for i in range(K)], dtype=torch.int32, device=q_weight.device - ) - - scale_mat = scales[current_g_idx] # (K, N) - scale_zeros_mat = scale_zeros[current_g_idx] # (K, N) - - # dequant: weight * scale - scale_zeros - dequantized_b = unpacked_q_weight.to(scales.dtype) * scale_mat - scale_zeros_mat - - return dequantized_b.reshape(K, N) - - -def torch_gptq_gemm( - a, b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, use_shuffle, bit -): - K, N = a.shape[1], b_q_weight.shape[1] - - b_dequant = torch_dequantize( - b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, use_shuffle, bit, K, N - ) - c = torch.matmul(a, b_dequant) - return c - - -def _test_gptq_gemm_once(M, N, K, bit, group_size, use_shuffle, dtype, device="cuda"): - - b_fp = torch.randn(K, N, dtype=dtype, device=device) - - assert K % group_size == 0, "K must be divisible by group_size" - num_groups = K // group_size - - if use_shuffle: - return - else: - g_idx = torch.tensor( - [i // group_size for i in range(K)], dtype=torch.int32, device=device - ) - b_shuffled = b_fp[g_idx] - - b_grouped = b_shuffled.reshape(num_groups, group_size, N) - - b_max = torch.max(b_grouped, dim=1, keepdim=True)[0] - b_min = torch.min(b_grouped, dim=1, keepdim=True)[0] - - scales = (b_max - b_min) / (2**bit - 1) - scales = scales.clamp(min=1e-6) - - zeros_float = (-b_min / scales).round() - - q_b = ( - (b_grouped / scales + zeros_float).round().clamp(0, 2**bit - 1).to(torch.uint8) - ) - - q_zeros_unpacked = zeros_float.to(torch.uint8) - 1 - - b_q_weight = pack_rows(q_b.reshape(K, N), bit, K, N) - - q_zeros_unpacked = q_zeros_unpacked.reshape(num_groups, N) - b_gptq_qzeros = pack_cols(q_zeros_unpacked, bit, num_groups, N) - b_gptq_scales = scales.squeeze(1) - - a = torch.randn(M, K, dtype=dtype, device=device) - - c_ref = torch_gptq_gemm( - a, b_q_weight, b_gptq_qzeros, b_gptq_scales, g_idx, use_shuffle, bit - ) - c_out = gptq_gemm( - a, b_q_weight, b_gptq_qzeros, b_gptq_scales, g_idx, use_shuffle, bit - ) - - rtol = 4e-2 - atol = 4e-2 - torch.testing.assert_close(c_ref, c_out, rtol=rtol, atol=atol) - print( - f"✅ Test passed: M={M}, N={N}, K={K}, bit={bit}, group_size={group_size}, use_shuffle={use_shuffle}, dtype={dtype}" - ) - - -@pytest.mark.parametrize("M", [1, 8, 128]) -@pytest.mark.parametrize("N", [2048, 4096]) -@pytest.mark.parametrize("K", [2048, 4096]) -@pytest.mark.parametrize("bit", [4]) -@pytest.mark.parametrize("group_size", [128]) -@pytest.mark.parametrize("use_shuffle", [False]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_gptq_gemm(M, N, K, bit, group_size, use_shuffle, dtype): - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") - _test_gptq_gemm_once(M, N, K, bit, group_size, use_shuffle, dtype, "cuda") - - -if __name__ == "__main__": - sys.exit(pytest.main([__file__, "-v"])) diff --git a/python/sglang/srt/arg_groups/attention_hook.py b/python/sglang/srt/arg_groups/attention_hook.py index d61d27d1a..e4f968285 100644 --- a/python/sglang/srt/arg_groups/attention_hook.py +++ b/python/sglang/srt/arg_groups/attention_hook.py @@ -9,7 +9,6 @@ from typing import Any from sglang.srt.arg_groups.overrides import ( _attention_backend_default, - _attention_backend_dual_chunk, _attention_backend_fa3_fp8_fallback, _attention_backend_platform_fallbacks, _cutedsl_prefill_backend_fill, @@ -209,23 +208,6 @@ def handle_attention_backend_compatibility(server_args: Any): # XPU platforms backends run_post_process_pass(server_args, _intel_xpu_page_constraint) - # Dual chunk flash attention backend - run_post_process_pass(server_args, _attention_backend_dual_chunk) - if resolved_view(server_args).attention_backend == "dual_chunk_flash_attn": - logger.warning( - "Mixed chunk and radix cache are disabled when using dual-chunk flash attention backend" - ) - declare_resolution( - server_args, - "_handle_attention_backend_compatibility", - enable_mixed_chunk=False, - ) - declare_resolution( - server_args, - "_handle_attention_backend_compatibility", - disable_radix_cache=True, - ) - def handle_linear_attn_backend(server_args: Any): cfg = resolving_view(server_args) diff --git a/python/sglang/srt/arg_groups/choices.py b/python/sglang/srt/arg_groups/choices.py index 954bc39d4..b5a663806 100644 --- a/python/sglang/srt/arg_groups/choices.py +++ b/python/sglang/srt/arg_groups/choices.py @@ -82,7 +82,6 @@ ATTENTION_BACKEND_CHOICES = [ "dsv4", "compressed", # Deprecated alias for "dsv4" # NVIDIA specific - "cutlass_mla", "fa3", "fa4", "flashinfer", @@ -91,7 +90,6 @@ ATTENTION_BACKEND_CHOICES = [ "cutedsl_mla", "tokenspeed_mla", "trtllm_mha", - "dual_chunk_flash_attn", "hpc_ops", # HPC-Ops (https://github.com/Tencent/hpc-ops), Hopper (SM90) only, requires --page-size 64 "minicpm_flashattn", "minicpm_flashinfer", @@ -235,7 +233,6 @@ CHUNKED_PREFIX_CACHE_SUPPORTED_ATTENTION_BACKENDS = [ "fa4", "flashmla", "cutedsl_mla", - "cutlass_mla", "trtllm_mla", "tokenspeed_mla", ] diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index c1d07e4b2..d872310fc 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -66,7 +66,6 @@ def handle_kv4_compatibility(server_args: Any) -> None: if prefill_backend == "fa4": if uses_mla: # FA4 + MLA KV4_FA4_MLA_BACKEND_CHOICES = [ - "cutlass_mla", "flashinfer", "trtllm_mla", ] @@ -87,7 +86,6 @@ def handle_kv4_compatibility(server_args: Any) -> None: else: if uses_mla: # !FA4 + MLA KV4_ATTENTION_MLA_BACKEND_CHOICES = [ - "cutlass_mla", "flashinfer", "trtllm_mla", ] @@ -371,7 +369,7 @@ def handle_page_major_kv_layout(server_args: Any): # Allow-list. Every backend below reads through the translator, so what # gates one is only whether its kernels can address the per-layer views: # * MLA models: the full paged MLA family, incl. flashmla (ps=64 - # snap). cutlass_mla stays rejected (never exercised). + # snap). # * MHA/SWA models: fa3 / fa4 / flashinfer / trtllm_mha alongside # Triton. fa4 is the fa3 class. # * Without the unified pool, plain page-major stays Triton-only. diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 4ece0f599..b73d49583 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -1153,14 +1153,6 @@ def _mla_backend_page_constraints(view: Any) -> dict: "FlashMLA only supports a page_size of 64, change page_size to 64." ) page_size = 64 - if ( - view.attention_backend == "cutlass_mla" - or view.decode_attention_backend == "cutlass_mla" - ): - logger.warning( - "Cutlass MLA only supports a page_size of 128, change page_size to 128." - ) - page_size = 128 if ( view.attention_backend == "trtllm_mla" or view.decode_attention_backend == "trtllm_mla" @@ -1369,23 +1361,6 @@ def _intel_xpu_page_constraint(view: Any) -> dict: return {} -@register_post_process -def _attention_backend_dual_chunk(view: Any) -> dict: - if ( - getattr(model_config_of(view).hf_config, "dual_chunk_attention_config", None) - is not None - ): - if view.attention_backend is None: - logger.info("Dual chunk attention is turned on by default.") - return {"attention_backend": "dual_chunk_flash_attn"} - elif view.attention_backend != "dual_chunk_flash_attn": - raise ValueError( - "Dual chunk attention is enabled, but attention backend is set to " - f"{view.attention_backend}. Please set it to 'dual_chunk_flash_attn'." - ) - return {} - - @register_post_process def _page_size_default(view: Any) -> dict: if view.page_size is not None: diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 24fd98d4b..f433f4141 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -981,7 +981,7 @@ def _handle_eagle_family(server_args: ServerArgs) -> None: # topk > 1 + page_size > 1 needs the two-pass cascade draft-decode (shared prefix # pass + per-branch expand pass with prefix-tail dup). Only these backends implement - # it; flashmla / trtllm_mla / cutlass_mla can't express the per-branch tree, so reject. + # it; flashmla / trtllm_mla can't express the per-branch tree, so reject. _PAGE_TREE_SPEC_BACKENDS = ("flashinfer", "fa3", "triton") view = resolved_view(server_args) if ( diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index d01c5d318..c99510301 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1657,7 +1657,6 @@ class ModelConfig: supported_quantization = [*QUANTIZATION_METHODS] rocm_supported_quantization = [ "awq", - "gptq", "fp8", "compressed_tensors", "compressed-tensors", diff --git a/python/sglang/srt/hardware_backend/gpu/quantization/awq_kernels.py b/python/sglang/srt/hardware_backend/gpu/quantization/awq_kernels.py index 91bffe663..ef4d7a0a7 100644 --- a/python/sglang/srt/hardware_backend/gpu/quantization/awq_kernels.py +++ b/python/sglang/srt/hardware_backend/gpu/quantization/awq_kernels.py @@ -68,10 +68,7 @@ else: awq_dequantize_triton as awq_dequantize, ) except ImportError: - try: - from sgl_kernel import awq_dequantize - except ImportError: - pass + pass _, scalar_types = get_scalar_types() diff --git a/python/sglang/srt/hardware_backend/gpu/quantization/gptq_kernels.py b/python/sglang/srt/hardware_backend/gpu/quantization/gptq_kernels.py index 580a4c22a..edf4e25b9 100644 --- a/python/sglang/srt/hardware_backend/gpu/quantization/gptq_kernels.py +++ b/python/sglang/srt/hardware_backend/gpu/quantization/gptq_kernels.py @@ -40,13 +40,9 @@ def _unsupported_kernel(*args, **kwargs): raise RuntimeError("GPTQ CUDA kernels are unavailable on the current platform.") -gptq_gemm = _unsupported_kernel gptq_marlin_repack = _unsupported_kernel -gptq_shuffle = _unsupported_kernel try: - from sgl_kernel import gptq_gemm, gptq_shuffle - from sglang.kernels.ops.quantization.gptq_marlin_repack import gptq_marlin_repack except Exception: pass @@ -82,52 +78,6 @@ def gptq_marlin_moe_repack( return output -class GPTQLinearKernel: - def __init__(self, quant_config: Optional[QuantizationConfig] = None): - self.quant_config = quant_config - self.use_shuffle = True - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - # for torch.compile - layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False) - layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False) - layer.g_idx = torch.nn.Parameter(layer.g_idx.data, requires_grad=False) - layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False) - - # exllama needs to shuffle the weight after the weight is loaded - # here we do the shuffle on first forward pass - if self.use_shuffle: - if self.quant_config.desc_act: - layer.g_idx.data = torch.argsort(layer.g_idx).to(torch.int) - else: - layer.g_idx.data = torch.empty( - (0,), dtype=torch.int, device=layer.g_idx.device - ) - gptq_shuffle(layer.qweight, layer.g_idx, self.quant_config.weight_bits) - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - out_shape = x.shape[:-1] + (layer.qweight.shape[-1],) - reshaped_x = x.reshape(-1, x.shape[-1]) - - output = gptq_gemm( - reshaped_x, - layer.qweight, - layer.qzeros, - layer.scales, - layer.g_idx, - self.use_shuffle, - self.quant_config.weight_bits, - ) - if bias is not None: - output.add_(bias) - return output.reshape(out_shape) - - class GPTQMarlinLinearKernel: def __init__(self, quant_config: Optional[QuantizationConfig] = None): self.quant_config = quant_config diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 30ff4f228..80dd33df8 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -259,13 +259,6 @@ def create_flashattention_v4_backend(runner): return FlashAttentionBackend(runner, fa_impl_ver=4) -@register_attention_backend("cutlass_mla") -def create_cutlass_mla_backend(runner): - from sglang.srt.layers.attention.cutlass_mla_backend import CutlassMLABackend - - return CutlassMLABackend(runner) - - @register_attention_backend("trtllm_mha") def create_trtllm_mha_backend(runner): if runner.use_mla_backend: @@ -299,15 +292,6 @@ def create_intel_amx_backend(runner): return IntelAMXAttnBackend(runner) -@register_attention_backend("dual_chunk_flash_attn") -def create_dual_chunk_flash_attn_backend(runner): - from sglang.srt.layers.attention.dual_chunk_flashattention_backend import ( - DualChunkFlashAttentionBackend, - ) - - return DualChunkFlashAttentionBackend(runner) - - def attn_backend_wrapper_for_draft_extend( runner: "ModelRunner", full_attn_backend: "AttentionBackend" ): diff --git a/python/sglang/srt/layers/attention/cutlass_mla_backend.py b/python/sglang/srt/layers/attention/cutlass_mla_backend.py deleted file mode 100644 index 410019ce2..000000000 --- a/python/sglang/srt/layers/attention/cutlass_mla_backend.py +++ /dev/null @@ -1,250 +0,0 @@ -from __future__ import annotations - -from sglang.srt.runtime_context import get_parallel - -""" -Support attention backend for Cutlass MLA. - -""" - -from dataclasses import dataclass -from typing import TYPE_CHECKING, Optional, Union - -import torch -import triton - -from sglang.kernels.ops.attention.utils import ( - create_flashmla_kv_indices_triton, - get_num_kv_index_blocks_flashmla, -) -from sglang.srt.layers.attention.flashinfer_mla_backend import FlashInferMLAAttnBackend -from sglang.srt.model_executor.forward_batch_info import ForwardBatch -from sglang.srt.utils import is_cuda - -if TYPE_CHECKING: - from sglang.srt.layers.radix_attention import RadixAttention - from sglang.srt.model_executor.model_runner import ModelRunner - -_is_cuda = is_cuda() -if _is_cuda: - from sgl_kernel import cutlass_mla_decode, cutlass_mla_get_workspace_size - - -# Cutlass MLA only supports pagesize=128 -PAGE_SIZE = 128 - - -@dataclass -class CutlassMLADecodeMetadata: - workspace: Optional[torch.Tensor] = None - block_kv_indices: Optional[torch.Tensor] = None - - def __init__( - self, - workspace: Optional[torch.Tensor] = None, - block_kv_indices: Optional[torch.Tensor] = None, - ): - self.workspace = workspace - self.block_kv_indices = block_kv_indices - - -class CutlassMLABackend(FlashInferMLAAttnBackend): - """Cutlass attention kernels.""" - - def __init__( - self, - model_runner: ModelRunner, - skip_prefill: bool = False, - kv_indptr_buf: Optional[torch.Tensor] = None, - kv_last_page_len_buf: Optional[torch.Tensor] = None, - ): - super().__init__( - model_runner, skip_prefill, kv_indptr_buf, kv_last_page_len_buf - ) - - self.num_q_heads = ( - model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size - ) - self.num_kv_heads = model_runner.model_config.get_num_kv_heads( - get_parallel().attn_tp_size - ) - self.req_to_token = model_runner.req_to_token_pool.req_to_token - self.num_local_heads = ( - model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size - ) - self.forward_metadata: Union[CutlassMLADecodeMetadata] = None - self.kv_lora_rank = model_runner.model_config.kv_lora_rank - self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim - self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim - self.v_head_dim = model_runner.model_config.v_head_dim - self.data_type = model_runner.kv_cache_dtype - self.q_data_type = model_runner.dtype - self.kv_cache_dim = self.kv_lora_rank + self.qk_rope_head_dim - - def init_forward_metadata_out_graph( - self, - forward_batch: ForwardBatch, - in_capture: bool = False, - ): - bs = forward_batch.batch_size - forward_mode = forward_batch.forward_mode - spec_info = forward_batch.spec_info - - if forward_mode.is_decode_or_idle() and spec_info is None: - create_flashmla_kv_indices_triton[ - ( - bs, - get_num_kv_index_blocks_flashmla( - self.cuda_graph_kv_indices.stride(0), PAGE_SIZE - ), - ) - ]( - self.req_to_token, - forward_batch.req_pool_indices[:bs], - forward_batch.seq_lens[:bs], - None, - self.cuda_graph_kv_indices, - self.req_to_token.stride(0), - self.cuda_graph_kv_indices.stride(0), - PAGED_SIZE=PAGE_SIZE, - ) - if in_capture: - max_seqlen_pad = self.cuda_graph_kv_indices.shape[1] - self.forward_metadata = CutlassMLADecodeMetadata( - self.cuda_graph_mla_workspace, - self.cuda_graph_kv_indices[:bs, :max_seqlen_pad], - ) - else: - super().init_forward_metadata_out_graph( - forward_batch, in_capture=in_capture - ) - - def init_forward_metadata(self, forward_batch: ForwardBatch): - - bs = forward_batch.batch_size - spec_info = forward_batch.spec_info - if forward_batch.forward_mode.is_decode_or_idle(): - if spec_info is None: - max_seqlen_pad = triton.cdiv( - forward_batch.seq_lens_cpu.max().item(), PAGE_SIZE - ) - block_kv_indices = torch.full( - (bs, max_seqlen_pad), - -1, - dtype=torch.int32, - device=forward_batch.seq_lens.device, - ) - create_flashmla_kv_indices_triton[ - (bs, get_num_kv_index_blocks_flashmla(max_seqlen_pad, PAGE_SIZE)) - ]( - self.req_to_token, - forward_batch.req_pool_indices, - forward_batch.seq_lens, - None, - block_kv_indices, - self.req_to_token.stride(0), - max_seqlen_pad, - PAGED_SIZE=PAGE_SIZE, - ) - workspace_size = cutlass_mla_get_workspace_size( - max_seqlen_pad * PAGE_SIZE, bs, num_kv_splits=1 - ) - workspace = torch.empty( - workspace_size, device="cuda", dtype=torch.uint8 - ) - self.forward_metadata = CutlassMLADecodeMetadata( - workspace, - block_kv_indices, - ) - else: - super().init_forward_metadata(forward_batch) - else: - super().init_forward_metadata(forward_batch) - - def init_cuda_graph_state( - self, - max_bs: int, - max_num_tokens: int, - block_kv_indices: Optional[torch.Tensor] = None, - ): - if block_kv_indices is None: - cuda_graph_kv_indices = torch.full( - (max_bs, (self.max_context_len + PAGE_SIZE) // PAGE_SIZE), - 1, - dtype=torch.int32, - device="cuda", - ) - else: - cuda_graph_kv_indices = block_kv_indices - - workspace_size = cutlass_mla_get_workspace_size( - cuda_graph_kv_indices.shape[1] * PAGE_SIZE, max_bs, num_kv_splits=1 - ) - self.cuda_graph_mla_workspace = torch.empty( - workspace_size, device="cuda", dtype=torch.uint8 - ) - self.cuda_graph_kv_indices = cuda_graph_kv_indices - - def get_cuda_graph_seq_len_fill_value(self): - return 1 - - def forward_decode( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - layer: RadixAttention, - forward_batch: ForwardBatch, - save_kv_cache: bool = True, - # For multi-head latent attention - q_rope: Optional[torch.Tensor] = None, - k_rope: Optional[torch.Tensor] = None, - ): - cache_loc = forward_batch.out_cache_loc - - if k is not None: - assert v is not None - if save_kv_cache: - if k_rope is not None: - self.token_to_kv_pool.set_mla_kv_buffer( - layer, - cache_loc, - k, - k_rope, - ) - else: - self.token_to_kv_pool.set_kv_buffer( - layer, - cache_loc, - k, - v, - ) - - # Reshape inputs - if q_rope is not None: - q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) - q_rope = q_rope.view( - -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim - ) - else: - reshaped_q = q.view(-1, layer.tp_q_head_num, layer.head_dim) - q_nope = reshaped_q[:, :, : layer.v_head_dim] - q_rope = reshaped_q[:, :, layer.v_head_dim :] - - q_nope = q_nope.to(self.q_data_type) - q_rope = q_rope.to(self.q_data_type) - - k_cache = self.token_to_kv_pool.get_key_buffer(layer.layer_id) - - o = cutlass_mla_decode( - q_nope=q_nope, - q_pe=q_rope, - kv_c_and_k_pe_cache=k_cache.view(-1, PAGE_SIZE, self.kv_cache_dim), - seq_lens=forward_batch.seq_lens.to(torch.int32), - page_table=self.forward_metadata.block_kv_indices, - workspace=self.forward_metadata.workspace, - sm_scale=layer.scaling, - num_kv_splits=1, - ) - - return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) diff --git a/python/sglang/srt/layers/attention/dual_chunk_flashattention_backend.py b/python/sglang/srt/layers/attention/dual_chunk_flashattention_backend.py deleted file mode 100644 index 83c22746e..000000000 --- a/python/sglang/srt/layers/attention/dual_chunk_flashattention_backend.py +++ /dev/null @@ -1,1711 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Attention layer with Dual chunk flash attention and sparse attention.""" - -import functools -import logging -import math -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple - -import torch -import torch.nn.functional as F -from sgl_kernel.sparse_flash_attn import ( - convert_vertical_slash_indexes, - convert_vertical_slash_indexes_mergehead, - sparse_attn_func, -) - -from sglang.kernels.ops.attention.flash_attention import ( - flash_attn_varlen_func, - flash_attn_with_kvcache, -) -from sglang.srt.layers.attention.base_attn_backend import AttentionBackend -from sglang.srt.layers.attention.flashattention_backend import ( - FlashAttentionMetadata, -) -from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode -from sglang.srt.runtime_context import get_parallel - -if TYPE_CHECKING: - from sglang.srt.layers.radix_attention import RadixAttention - from sglang.srt.model_executor.model_runner import ModelRunner - - -logger = logging.getLogger(__name__) - - -@dataclass -class DualChunkFlashAttentionMetadata: - """Metadata for FlashAttentionBackend. - - NOTE: Any python object stored here is not updated when it is - cuda-graph replayed. If you have values that need to be changed - dynamically, it should be stored in tensor. The tensor has to be - updated from `CUDAGraphRunner.forward` API. - """ - - # (batch_size,). The sequence length per sequence. Sequence length means - # the computed tokens + new tokens None if it is a decoding. - seq_lens: Optional[List[int]] = None - # seq_lens stored as a tensor. - seq_lens_tensor: Optional[torch.Tensor] = None - # Maximum sequence length among prefill batch. 0 if there are decoding - # requests only. - max_seq_len: int = None - - # (batch_size,). The orig sequence length per sequence. - orig_seq_lens: Optional[List[int]] = None - - # orig_seq_lens stored as a tensor. - orig_seq_lens_tensor: Optional[torch.Tensor] = None - - # Block addresses per sequence. (Seq id -> list of physical block) - # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks - # in the kv cache. Each block can contain up to block_size tokens. - # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph - # captured. - block_tables: Optional[torch.Tensor] = None - - # (batch_size + 1,). The cumulative subquery lengths of the sequences in - # the batch, used to index into subquery. E.g., if the subquery length - # is [4, 6], it is [0, 4, 10]. - query_start_loc: Optional[torch.Tensor] = None - # (batch_size + 1,). The cumulative sequence lengths of the sequences in - # the batch, used to index into sequence. E.g., if the sequence length is - # [4, 6], it is [0, 4, 10]. - seq_start_loc: Optional[torch.Tensor] = None - - # Length scaling factor - scaling_factor: Optional[torch.Tensor] = None - - # (batch_size,). Sequence lengths for intra attention. - seq_lens_intra: Optional[torch.Tensor] = None - - # Max sequence length for intra attention. - max_seq_len_intra: Optional[int] = None - - # (batch_size, num_blocks). Block table for intra attention. - block_tables_intra: Optional[torch.Tensor] = None - - # (batch_size,). Sequence lengths for succ attention. - seq_lens_succ: Optional[torch.Tensor] = None - - # Max sequence length for succ attention. - max_seq_len_succ: Optional[int] = None - - # (batch_size, num_blocks). Block table for succ attention. - block_tables_succ: Optional[torch.Tensor] = None - - # (batch_size,). Sequence lengths for inter attention. - seq_lens_inter: Optional[torch.Tensor] = None - - # Max sequence length for inter attention. - max_seq_len_inter: Optional[int] = None - - -class DualChunkFlashAttentionBackend(AttentionBackend): - def __init__( - self, - model_runner: "ModelRunner", - ) -> None: - self.forward_metadata: FlashAttentionMetadata = None - self.device = model_runner.device - self.max_context_len = model_runner.model_config.context_len - self.num_heads = model_runner.model_config.get_num_attention_heads( - get_parallel().tp_size - ) - self.num_kv_heads = model_runner.model_config.get_num_kv_heads( - get_parallel().tp_size - ) - self.head_size = model_runner.model_config.head_dim - - # Pool refs — captured at construction so they survive deletion of the - # corresponding ForwardBatch fields. - self.req_to_token_pool = model_runner.req_to_token_pool - self.token_to_kv_pool = model_runner.token_to_kv_pool - self.req_to_token = model_runner.req_to_token_pool.req_to_token - self.kv_cache_dtype = model_runner.kv_cache_dtype - - self.kv_cache_dtype_str = model_runner.kv_cache_dtype_str - self.page_size = model_runner.page_size - - assert self.num_heads % self.num_kv_heads == 0 - self.num_queries_per_kv = self.num_heads // self.num_kv_heads - - dual_chunk_attention_config = getattr( - model_runner.model_config.hf_config, "dual_chunk_attention_config", None - ) - assert dual_chunk_attention_config is not None - self.chunk_size = dual_chunk_attention_config.get("chunk_size", 8192) - self.local_size = dual_chunk_attention_config.get("local_size", 1024) - self.original_max_position_embeddings = dual_chunk_attention_config.get( - "original_max_position_embeddings", 0 - ) - self.sparse_attention_config = dual_chunk_attention_config.get( - "sparse_attention_config", None - ) - if not self.sparse_attention_config: - logger.warning_once( - "Sparse attention will not be enabled as " - "sparse attention config is not provided." - ) - self.sparse_attention_enabled = dual_chunk_attention_config.get( - "sparse_attention_enabled", self.sparse_attention_config is not None - ) - self.sparse_attention_threshold = dual_chunk_attention_config.get( - "sparse_attention_threshold", 32768 - ) - self.sparse_attention_last_q = dual_chunk_attention_config.get( - "sparse_attention_last_q", 64 - ) - self.dual_chunk_attention_config = dual_chunk_attention_config - - if self.sparse_attention_enabled: - self.arange = torch.arange(self.sparse_attention_last_q, device="cuda") - self.last_q_mask = ( - self.arange[None, None, :, None] >= self.arange[None, None, None, :] - ) - - @functools.lru_cache() - def get_sparse_attention_config(self, layer_idx) -> List[Dict[str, Any]]: - layer_sparse_attention_config = { - int(i): j for i, j in self.sparse_attention_config[layer_idx].items() - } - start_head = self.num_heads * get_parallel().tp_rank - end_head = start_head + self.num_heads - return [layer_sparse_attention_config[i] for i in range(start_head, end_head)] - - def init_forward_metadata_out_graph( - self, - forward_batch: ForwardBatch, - in_capture: bool = False, - ): - bs = forward_batch.batch_size - req_pool_indices = forward_batch.req_pool_indices - seq_lens = forward_batch.seq_lens - forward_mode = forward_batch.forward_mode - - if in_capture: - self._bind_metadata_buffers(bs, req_pool_indices, forward_mode) - - self._apply_cuda_graph_metadata( - bs=bs, - req_pool_indices=req_pool_indices, - seq_lens=seq_lens, - forward_mode=forward_mode, - ) - - if in_capture and forward_mode.is_decode_or_idle(): - # Restore max_seq_len scalars — replay sets actual values but CUDA - # graph needs the safe upper bound baked in at capture time. - md = self.forward_metadata - md.max_seq_len = self.max_context_len - md.max_seq_len_intra = self.max_context_len - md.max_seq_len_succ = self.max_context_len - md.max_seq_len_inter = self.max_context_len - - def init_forward_metadata(self, forward_batch: ForwardBatch): - """Initialize forward metadata hence all layers in the forward pass can reuse it.""" - - forward_mode: ForwardMode = forward_batch.forward_mode - assert forward_mode.is_prefill() or forward_mode.is_decode() - batch_size = forward_batch.batch_size - - metadata = DualChunkFlashAttentionMetadata() - metadata.seq_lens_tensor = forward_batch.seq_lens.to(torch.int32) - metadata.seq_lens = forward_batch.seq_lens.tolist() - metadata.max_seq_len = forward_batch.seq_lens.max().item() - - metadata.orig_seq_lens_tensor = forward_batch.orig_seq_lens - metadata.orig_seq_lens = forward_batch.orig_seq_lens.tolist() - - metadata.block_tables = self.req_to_token_pool.req_to_token[ - forward_batch.req_pool_indices, : metadata.max_seq_len - ] - # Convert the block table to a strided format. - if self.page_size > 1: - strided_indices = torch.arange( - 0, metadata.block_tables.shape[1], self.page_size, device=self.device - ) - metadata.block_tables = ( - metadata.block_tables[:, strided_indices] // self.page_size - ) - - metadata.query_start_loc = torch.zeros( - batch_size + 1, dtype=torch.int32, device=metadata.seq_lens_tensor.device - ) - if forward_mode.is_prefill(): - metadata.query_start_loc[1:] = torch.cumsum( - forward_batch.extend_seq_lens.to(torch.int32), dim=0, dtype=torch.int32 - ) - else: - metadata.query_start_loc[1:] = torch.cumsum( - torch.arange( - batch_size, - dtype=metadata.query_start_loc.dtype, - device=metadata.query_start_loc.device, - ), - dim=0, - dtype=torch.int32, - ) - metadata.seq_start_loc = torch.zeros( - batch_size + 1, dtype=torch.int32, device=metadata.seq_lens_tensor.device - ) - metadata.seq_start_loc[1:] = torch.cumsum( - metadata.seq_lens_tensor, dim=0, dtype=torch.int32 - ) - - if self.original_max_position_embeddings > 0: - if forward_mode.is_prefill(): - metadata.scaling_factor = ( - 0.1 - * torch.log( - metadata.orig_seq_lens_tensor - / self.original_max_position_embeddings - ) - + 1.0 - ).clip(min=1) - else: - metadata.scaling_factor = ( - 0.1 - * torch.log( - metadata.orig_seq_lens_tensor - / self.original_max_position_embeddings - ) - + 1.0 - ).clip(min=1) - - if forward_mode.is_decode(): - cache_seq_lens = metadata.orig_seq_lens_tensor - - chunk_len = self.chunk_size - self.local_size - chunk_num_curr = (cache_seq_lens - 1) // chunk_len - - seq_lens_intra = cache_seq_lens - chunk_num_curr * chunk_len - max_seq_len_intra = seq_lens_intra.max().item() - metadata.seq_lens_intra = seq_lens_intra - metadata.max_seq_len_intra = max_seq_len_intra - - block_tables_intra = torch.zeros( - batch_size, - (max_seq_len_intra - 1) // self.page_size + 1, - dtype=metadata.block_tables.dtype, - device=metadata.block_tables.device, - ) - for i in range(batch_size): - st = chunk_num_curr[i] * chunk_len // self.page_size - ed = min( - st + (max_seq_len_intra - 1) // self.page_size + 1, - (cache_seq_lens[i] - 1) // self.page_size + 1, - ) - block_tables_intra[i, : ed - st] = metadata.block_tables[i, st:ed] - metadata.block_tables_intra = block_tables_intra - - metadata.seq_lens_succ = ( - chunk_num_curr - (chunk_num_curr - 1).clip(min=0) - ) * chunk_len - metadata.max_seq_len_succ = metadata.seq_lens_succ.max().item() - if metadata.max_seq_len_succ: - block_tables_succ = torch.zeros( - batch_size, - (metadata.max_seq_len_succ - 1) // self.page_size + 1, - dtype=metadata.block_tables.dtype, - device=metadata.block_tables.device, - ) - for i in range(batch_size): - start = ( - (chunk_num_curr[i] - 1).clip(min=0) - * chunk_len - // self.page_size - ) - end = min( - start + (metadata.max_seq_len_succ - 1) // self.page_size + 1, - (cache_seq_lens[i] - 1) // self.page_size + 1, - ) - block_tables_succ[i, : end - start] = metadata.block_tables[ - i, start:end - ] - metadata.block_tables_succ = block_tables_succ - - metadata.seq_lens_inter = (chunk_num_curr - 1).clip(min=0) * chunk_len - metadata.max_seq_len_inter = metadata.seq_lens_inter.max().item() - - self.forward_metadata = metadata - - def forward_extend( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - layer: "RadixAttention", - forward_batch: ForwardBatch, - save_kv_cache=True, - ): - # Use precomputed metadata across all layers - metadata = self.forward_metadata - - ( - query, - query_succ, - query_inter, - query_succ_critical, - query_inter_critical, - ) = torch.split(q, q.shape[-1] // 5, dim=-1) - - # Reshape the query, key, and value tensors. - query = query.view(-1, self.num_heads, self.head_size) - query_succ = query_succ.view(-1, self.num_heads, self.head_size) - query_inter = query_inter.view(-1, self.num_heads, self.head_size) - query_succ_critical = query_succ_critical.view( - -1, self.num_heads, self.head_size - ) - query_inter_critical = query_inter_critical.view( - -1, self.num_heads, self.head_size - ) - key = k.view(-1, self.num_kv_heads, self.head_size) - value = v.view(-1, self.num_kv_heads, self.head_size) - - # apply DCA scaling - if self.original_max_position_embeddings > 0: - assert metadata.scaling_factor is not None - assert metadata.query_start_loc is not None - assert metadata.orig_seq_lens is not None - current_start = 0 - query_start_loc_cpu = metadata.query_start_loc.cpu() - for i in range(len(metadata.orig_seq_lens)): - current_end = ( - current_start - + (query_start_loc_cpu[i + 1] - query_start_loc_cpu[i]).item() - ) - key[current_start:current_end].mul_(metadata.scaling_factor[i]) - current_start = current_end - assert current_end <= self.max_context_len - - # Do multi-head attention - key_cache, value_cache = self.token_to_kv_pool.get_kv_buffer(layer.layer_id) - key_cache = key_cache.view( - -1, self.page_size, layer.tp_k_head_num, layer.head_dim - ) - value_cache = value_cache.view( - -1, self.page_size, layer.tp_v_head_num, layer.head_dim - ) - - if key is not None and value is not None: - if save_kv_cache: - self.token_to_kv_pool.set_kv_buffer( - layer, - forward_batch.out_cache_loc, - key, - value, - layer.k_scale, - layer.v_scale, - ) - - if not save_kv_cache: - # profile run - o = flash_attn_varlen_func( - q=query, - k=key, - v=value, - cu_seqlens_q=metadata.seq_start_loc, - cu_seqlens_k=metadata.seq_start_loc, - max_seqlen_q=metadata.max_seq_len, - max_seqlen_k=metadata.max_seq_len, - softmax_scale=layer.scaling, - causal=True, - ) - else: - # prefill/chunked-prefill - # get per layer sparse attention config - if self.sparse_attention_enabled: - self.layer_sparse_attention_config = self.get_sparse_attention_config( - layer.layer_id - ) - assert metadata.orig_seq_lens is not None - o = self._dual_chunk_flash_attn_prefill( - q=query, - q_succ=query_succ, - q_inter=query_inter, - q_succ_critical=query_succ_critical, - q_inter_critical=query_inter_critical, - k=key_cache, - v=value_cache, - cu_seqlens_q=metadata.query_start_loc, - cu_seqlens_k=metadata.seq_start_loc, - orig_seq_lens=metadata.orig_seq_lens, - scaling_factor=metadata.scaling_factor, - softmax_scale=layer.scaling, - causal=True, - window_size=(-1, -1), - block_table=metadata.block_tables, - chunk_size=self.chunk_size, - local_size=self.local_size, - ) - return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) - - def forward_decode( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - layer: "RadixAttention", - forward_batch: ForwardBatch, - save_kv_cache=True, - ) -> torch.Tensor: - # Use precomputed metadata across all layers - metadata = self.forward_metadata - - ( - query, - query_succ, - query_inter, - query_succ_critical, - query_inter_critical, - ) = torch.split(q, q.shape[-1] // 5, dim=-1) - - # Reshape the query, key, and value tensors. - query = query.view(-1, self.num_heads, self.head_size) - query_succ = query_succ.view(-1, self.num_heads, self.head_size) - query_inter = query_inter.view(-1, self.num_heads, self.head_size) - query_succ_critical = query_succ_critical.view( - -1, self.num_heads, self.head_size - ) - query_inter_critical = query_inter_critical.view( - -1, self.num_heads, self.head_size - ) - key = k.view(-1, self.num_kv_heads, self.head_size) - value = v.view(-1, self.num_kv_heads, self.head_size) - - key_cache, value_cache = self.token_to_kv_pool.get_kv_buffer(layer.layer_id) - key_cache = key_cache.view( - -1, self.page_size, layer.tp_k_head_num, layer.head_dim - ) - value_cache = value_cache.view( - -1, self.page_size, layer.tp_v_head_num, layer.head_dim - ) - - if key is not None and value is not None: - if save_kv_cache: - self.token_to_kv_pool.set_kv_buffer( - layer, - forward_batch.out_cache_loc, - key, - value, - layer.k_scale, - layer.v_scale, - ) - - # apply DCA scaling - if self.original_max_position_embeddings > 0: - assert metadata.scaling_factor is not None - scaling_factor = metadata.scaling_factor - key.mul_(scaling_factor.unsqueeze(-1).unsqueeze(-1)) - - o = self._dual_chunk_flash_attn_decoding( - query.unsqueeze(1), - query_succ.unsqueeze(1), - query_inter.unsqueeze(1), - key_cache, - value_cache, - block_table=metadata.block_tables, - cache_seqlens=metadata.seq_lens_tensor, - softmax_scale=layer.scaling, - causal=True, - chunk_size=self.chunk_size, - local_size=self.local_size, - original_max_position_embeddings=self.original_max_position_embeddings, - decode_meta=metadata, - ).squeeze(1) - return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) - - def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): - """Initialize CUDA graph state for the attention backend. - - Args: - max_bs (int): Maximum batch size to support in CUDA graphs - - This creates fixed-size tensors that will be reused during CUDA graph replay - to avoid memory allocations. - """ - self.decode_metadata = { - "seq_lens_tensor": torch.zeros( - max_bs, dtype=torch.int32, device=self.device - ), - "orig_seq_lens_tensor": torch.zeros( - max_bs, dtype=torch.int32, device=self.device - ), - "scaling_factor": torch.zeros( - max_bs, dtype=torch.float32, device=self.device - ), - "block_tables": torch.zeros( - max_bs, - (self.max_context_len - 1) // self.page_size + 1, - dtype=torch.int32, - device=self.device, - ), - "block_tables_intra": torch.zeros( - max_bs, - (self.max_context_len - 1) // self.page_size + 1, - dtype=torch.int32, - device=self.device, - ), - "seq_lens_intra": torch.zeros( - max_bs, dtype=torch.int32, device=self.device - ), - "block_tables_succ": torch.zeros( - max_bs, - (self.max_context_len - 1) // self.page_size + 1, - dtype=torch.int32, - device=self.device, - ), - "seq_lens_succ": torch.zeros(max_bs, dtype=torch.int32, device=self.device), - "seq_lens_inter": torch.zeros( - max_bs, dtype=torch.int32, device=self.device - ), - } - - def _bind_metadata_buffers( - self, - bs: int, - req_pool_indices: torch.Tensor, - forward_mode: ForwardMode, - ): - """Allocate persistent metadata buffers for CUDA graph capture.""" - metadata = DualChunkFlashAttentionMetadata() - - if forward_mode.is_decode_or_idle(): - if self.original_max_position_embeddings > 0: - metadata.scaling_factor = self.decode_metadata["scaling_factor"][:bs] - - metadata.seq_lens_tensor = self.decode_metadata["seq_lens_tensor"][:bs] - metadata.orig_seq_lens_tensor = self.decode_metadata[ - "orig_seq_lens_tensor" - ][:bs] - metadata.max_seq_len = self.max_context_len - metadata.block_tables = self.decode_metadata["block_tables"][ - req_pool_indices, : - ] - - # intra - metadata.max_seq_len_intra = self.max_context_len - metadata.seq_lens_intra = self.decode_metadata["seq_lens_intra"][:bs] - - metadata.block_tables_intra = self.decode_metadata["block_tables_intra"][ - :bs, : - ] - - # succ - metadata.seq_lens_succ = self.decode_metadata["seq_lens_succ"][:bs] - metadata.max_seq_len_succ = self.max_context_len - - metadata.block_tables_succ = self.decode_metadata["block_tables_succ"][ - :bs, : - ] - - metadata.seq_lens_inter = self.decode_metadata["seq_lens_inter"][:bs] - metadata.max_seq_len_inter = self.max_context_len - - self.decode_metadata[bs] = metadata - - self.forward_metadata = metadata - - def _apply_cuda_graph_metadata( - self, - bs: int, - req_pool_indices: torch.Tensor, - seq_lens: torch.Tensor, - forward_mode: ForwardMode, - ): - """Shared capture+replay body for the cuda-graph init path. - - Public entry: :py:meth:`init_forward_metadata_out_graph`. - """ - assert forward_mode.is_decode() - seq_lens = seq_lens[:bs] - req_pool_indices = req_pool_indices[:bs] - metadata = self.decode_metadata[bs] - - metadata.seq_lens_tensor.copy_(seq_lens.to(torch.int32)) - metadata.seq_lens = seq_lens.tolist() - metadata.max_seq_len = seq_lens.max().item() - - metadata.orig_seq_lens_tensor.copy_(seq_lens) - metadata.orig_seq_lens = seq_lens.tolist() - - block_tables = self.req_to_token[req_pool_indices, : metadata.max_seq_len] - # Convert the block table to a strided format. - if self.page_size > 1: - strided_indices = torch.arange( - 0, block_tables.shape[1], self.page_size, device=self.device - ) - block_tables = block_tables[:, strided_indices] // self.page_size - metadata.block_tables.fill_(0) - metadata.block_tables[: block_tables.shape[0], : block_tables.shape[1]].copy_( - block_tables - ) - - if self.original_max_position_embeddings > 0: - scaling_factor = ( - 0.1 - * torch.log( - metadata.orig_seq_lens_tensor - / self.original_max_position_embeddings - ) - + 1.0 - ).clip(min=1) - metadata.scaling_factor.copy_(scaling_factor) - - cache_seq_lens = metadata.orig_seq_lens_tensor - - chunk_len = self.chunk_size - self.local_size - chunk_num_curr = (cache_seq_lens - 1) // chunk_len - - seq_lens_intra = cache_seq_lens - chunk_num_curr * chunk_len - max_seq_len_intra = seq_lens_intra.max().item() - metadata.seq_lens_intra.copy_(seq_lens_intra) - metadata.max_seq_len_intra = max_seq_len_intra - - metadata.block_tables_intra.fill_(0) - for i in range(bs): - st = chunk_num_curr[i] * chunk_len // self.page_size - ed = min( - st + (max_seq_len_intra - 1) // self.page_size + 1, - (cache_seq_lens[i] - 1) // self.page_size + 1, - ) - metadata.block_tables_intra[i, : ed - st] = metadata.block_tables[i, st:ed] - - seq_lens_succ = (chunk_num_curr - (chunk_num_curr - 1).clip(min=0)) * chunk_len - metadata.seq_lens_succ.copy_(seq_lens_succ) - metadata.max_seq_len_succ = metadata.seq_lens_succ.max().item() - if metadata.max_seq_len_succ: - metadata.block_tables_succ.fill_(0) - for i in range(bs): - start = ( - (chunk_num_curr[i] - 1).clip(min=0) * chunk_len // self.page_size - ) - end = min( - start + (metadata.max_seq_len_succ - 1) // self.page_size + 1, - (cache_seq_lens[i] - 1) // self.page_size + 1, - ) - metadata.block_tables_succ[i, : end - start] = metadata.block_tables[ - i, start:end - ] - - seq_lens_inter = (chunk_num_curr - 1).clip(min=0) * chunk_len - metadata.seq_lens_inter.copy_(seq_lens_inter) - metadata.max_seq_len_inter = metadata.seq_lens_inter.max().item() - - self.forward_metadata = metadata - - def get_cuda_graph_seq_len_fill_value(self): - """Get the fill value for sequence length in CUDA graph.""" - return 1 - - def _dual_chunk_flash_attn_prefill( - self, - q, - q_succ, - q_inter, - q_succ_critical, - q_inter_critical, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - orig_seq_lens: List[int], - scaling_factor: torch.Tensor, - softmax_scale: float, - causal: Optional[bool] = True, - window_size: Tuple[int, int] = (-1, -1), - block_table: Optional[torch.Tensor] = None, - chunk_size: int = 8192, - local_size: int = 1024, - ): - if not causal: - raise ValueError("Dual Chunk Attention does not support causal=False") - if window_size != (-1, -1): - raise ValueError("Dual Chunk Attention does not support window_size") - - cu_seqlens_q_cpu = cu_seqlens_q.cpu().tolist() - cu_seqlens_k_cpu = cu_seqlens_k.cpu().tolist() - all_outputs = [] - - for i in range(0, len(cu_seqlens_q_cpu) - 1): - qs = cu_seqlens_q_cpu[i] - qe = cu_seqlens_q_cpu[i : i + 2][-1] - ks = cu_seqlens_k_cpu[i] - ke = cu_seqlens_k_cpu[i : i + 2][-1] - - current_q = q[qs:qe] - current_q_succ = q_succ[qs:qe] - current_q_inter = q_inter[qs:qe] - current_q_succ_critical = q_succ_critical[qs:qe] - current_q_inter_critical = q_inter_critical[qs:qe] - - if block_table is None: - current_k = k[ks:ke] - current_v = v[ks:ke] - current_block_table = None - current_orig_seq_len = orig_seq_lens[i] - else: - current_block_table = block_table[i] - current_orig_seq_len = orig_seq_lens[i] - current_k = k - current_v = v - sparse_attn_enabled = ( - self.sparse_attention_enabled - and current_orig_seq_len > self.sparse_attention_threshold - ) - - if current_q.shape[0] == 0: - continue - - if current_k.shape[0] == 0: - all_outputs.append( - torch.zeros( - (current_q.shape[0], current_q.shape[1], v.shape[2]), - device=q.device, - dtype=q.dtype, - ) - ) - continue - - current_output = torch.empty_like(current_q) - group_size = int(current_q.size(-2) / current_k.size(-2)) - - if sparse_attn_enabled: - num_device_q_heads = current_q.size(-2) - heads_vertical_size = torch.empty( - size=(num_device_q_heads,), dtype=torch.int32 - ) - heads_slash_size = torch.empty( - size=(num_device_q_heads,), dtype=torch.int32 - ) - for head_id in range(current_q.size(-2)): - ( - ty, - vertical_size, - slash_size, - _, - ) = self.layer_sparse_attention_config[head_id] - assert ty == "vertical_and_slash", "only support slash mode" - - if vertical_size == 30: - vertical_size += 100 - heads_vertical_size[head_id] = vertical_size - heads_slash_size[head_id] = slash_size - - current_output = self._dual_chunk_flash_attn_prefill_func( - current_q, # allheads - current_q_succ, - current_q_inter, - current_q_succ_critical, - current_q_inter_critical, - current_k, - current_v, - current_block_table, - softmax_scale, - chunk_size, - local_size, - scaling_factor[i].item(), - ke - ks, - sparse_attn_enabled=sparse_attn_enabled, - heads_vertical_size=heads_vertical_size, - heads_slash_size=heads_slash_size, - group_size=group_size, - ) - else: - for head_id in range(current_q.size(-2)): - # (seq_len, num_heads, head_size) - current_q_head = current_q[:, head_id, :].unsqueeze(1) - current_q_succ_head = current_q_succ[:, head_id, :].unsqueeze(1) - current_q_inter_head = current_q_inter[:, head_id, :].unsqueeze(1) - current_q_succ_head_critical = current_q_succ_critical[ - :, head_id, : - ].unsqueeze(1) - current_q_inter_head_critical = current_q_inter_critical[ - :, head_id, : - ].unsqueeze(1) - if block_table is not None: - current_k_head = current_k[ - ..., head_id // group_size, : - ].unsqueeze(2) - current_v_head = current_v[ - ..., head_id // group_size, : - ].unsqueeze(2) - - else: - current_k_head = current_k[:, head_id, :].unsqueeze(1) - current_v_head = current_v[:, head_id, :].unsqueeze(1) - - current_out = self._dual_chunk_flash_attn_prefill_func( - current_q_head, - current_q_succ_head, - current_q_inter_head, - current_q_succ_head_critical, - current_q_inter_head_critical, - current_k_head, - current_v_head, - current_block_table, - softmax_scale, - chunk_size, - local_size, - scaling_factor[i].item(), - ke - ks, - sparse_attn_enabled=sparse_attn_enabled, - ) - current_output[:, head_id : head_id + 1, :] = current_out - all_outputs.append(current_output) - return torch.cat(all_outputs, dim=0) - - def _dual_chunk_flash_attn_prefill_func( - self, - q, - q_succ, - q_inter, - q_succ_critical, - q_inter_critical, - k, - v, - block_table, - softmax_scale: float, - chunk_size: int, - local_size: int, - scaling_factor: float, - k_length: int, - sparse_attn_enabled: Optional[bool] = True, - heads_vertical_size=None, - heads_slash_size=None, - group_size=None, - ): - flash_results = [] - chunk_len = chunk_size - local_size - - if block_table is not None: - block_size = v.shape[1] - if chunk_len % block_size != 0: - raise ValueError("chunk_len must be divisible by block_size.") - else: - block_size = 1 - - if self.original_max_position_embeddings > 0: - softmax_scale = softmax_scale * scaling_factor - - begin = k_length - q.shape[0] - while begin < k_length: - flash_per_chunk = [] - - prev_chunk_end_pos = (begin // chunk_len) * chunk_len - next_chunk_end_pos = prev_chunk_end_pos + chunk_len - end = min(next_chunk_end_pos, k_length) - qbegin = begin - (k_length - q.shape[0]) - qend = end - (k_length - q.shape[0]) - - qk_chunks = [] - q_states_intra = q[qbegin:qend] - # choose critical token - if block_table is not None: - block_tables_intra = _get_block( - block_table, block_size, prev_chunk_end_pos, end - ) - k_states_intra = k[block_tables_intra].view(-1, *k.shape[-2:])[ - : (end - prev_chunk_end_pos) - ] - v_states_intra = v[block_tables_intra].view(-1, *v.shape[-2:])[ - : (end - prev_chunk_end_pos) - ] - else: - block_tables_intra = None - k_states_intra = k[prev_chunk_end_pos:end] - v_states_intra = v[prev_chunk_end_pos:end] - - if sparse_attn_enabled: - last_q_size = min(qend - qbegin, self.sparse_attention_last_q) - _, num_device_k_heads, head_dim = k_states_intra.shape - k_states_intra = ( - k_states_intra.unsqueeze(2) - .repeat(1, 1, group_size, 1) - .reshape(-1, num_device_k_heads * group_size, head_dim) - ) - v_states_intra = ( - v_states_intra.unsqueeze(2) - .repeat(1, 1, group_size, 1) - .reshape(-1, num_device_k_heads * group_size, head_dim) - ) - qk_chunks.append( - (q_states_intra.transpose(0, 1)[:, -last_q_size:] * softmax_scale) - @ k_states_intra.permute(1, 2, 0) - ) - - if prev_chunk_end_pos - chunk_len >= 0: - q_states_succ = q_succ[qbegin:qend] - q_states_succ_critical = q_succ_critical[qbegin:qend] - if block_table is not None: - block_tables_succ = _get_block( - block_table, - block_size, - prev_chunk_end_pos - chunk_len, - prev_chunk_end_pos, - ) - k_states_succ = k[block_tables_succ].view(-1, *k.shape[-2:])[ - :chunk_len - ] - v_states_succ = v[block_tables_succ].view(-1, *v.shape[-2:])[ - :chunk_len - ] - else: - k_states_succ = k[ - prev_chunk_end_pos - chunk_len : prev_chunk_end_pos - ] - v_states_succ = v[ - prev_chunk_end_pos - chunk_len : prev_chunk_end_pos - ] - - if sparse_attn_enabled: - k_states_succ = ( - k_states_succ.unsqueeze(2) - .repeat(1, 1, group_size, 1) - .reshape(-1, num_device_k_heads * group_size, head_dim) - ) - v_states_succ = ( - v_states_succ.unsqueeze(2) - .repeat(1, 1, group_size, 1) - .reshape(-1, num_device_k_heads * group_size, head_dim) - ) - qk_chunks.append( - ( - q_states_succ_critical.transpose(0, 1)[:, -last_q_size:] - * softmax_scale - ) - @ k_states_succ.permute(1, 2, 0) - ) - - if prev_chunk_end_pos - chunk_len * 2 >= 0: - q_states_inter = q_inter[qbegin:qend] - q_states_inter_critical = q_inter_critical[qbegin:qend] - if block_table is not None: - block_tables_inter = _get_block( - block_table, block_size, 0, prev_chunk_end_pos - chunk_len - ) - k_states_inter = k[block_tables_inter].view(-1, *k.shape[-2:])[ - : (prev_chunk_end_pos - chunk_len) - ] - v_states_inter = v[block_tables_inter].view(-1, *v.shape[-2:])[ - : (prev_chunk_end_pos - chunk_len) - ] - else: - k_states_inter = k[: prev_chunk_end_pos - chunk_len] - v_states_inter = v[: prev_chunk_end_pos - chunk_len] - - if sparse_attn_enabled: - k_states_inter = ( - k_states_inter.unsqueeze(2) - .repeat(1, 1, group_size, 1) - .reshape(-1, num_device_k_heads * group_size, head_dim) - ) - v_states_inter = ( - v_states_inter.unsqueeze(2) - .repeat(1, 1, group_size, 1) - .reshape(-1, num_device_k_heads * group_size, head_dim) - ) - qk_chunks.append( - ( - q_states_inter_critical.transpose(0, 1)[:, -last_q_size:] - * softmax_scale - ) - @ k_states_inter.permute(1, 2, 0) - ) - - if sparse_attn_enabled: - reversed_qk = qk_chunks[::-1] - qk = torch.cat(reversed_qk, dim=-1) - - qk[:, :, -last_q_size:] = torch.where( - self.last_q_mask[..., -last_q_size:, -last_q_size:].to(qk.device), - qk[:, :, -last_q_size:], - -torch.inf, - ) - qk = F.softmax(qk, dim=-1, dtype=torch.float32) - - vertical = qk.sum(-2, keepdim=True) - vertical[..., :30] = torch.inf - - # Avoid sorting by using the min/max ints to fill the indexer - # buffers. - int32_max = torch.iinfo(torch.int32).max - int32_min = torch.iinfo(torch.int32).min - n_heads = qk.size()[0] - max_slash_topk = torch.max(heads_slash_size).item() - max_vertical_topk = torch.max(heads_vertical_size).item() - # store each head's slash topk, vertical topk - vertical = vertical.reshape((n_heads, -1)) - # prevent out of range when prompt size < max_vertical_topk - max_vertical_topk = min(vertical.shape[-1], max_vertical_topk) - vertical_topk_buffer = torch.topk( - vertical, max_vertical_topk, -1 - ).indices - slash_topk_buffer = torch.empty( - size=(n_heads, max_slash_topk), dtype=torch.int64, device=qk.device - ) - for head_i in range(n_heads): - # (nqheads=1, lastq, k_len) - head_score = qk[head_i : head_i + 1, :, :] - slash_scores = _sum_all_diagonal_matrix(head_score) - if head_score.size(1) != 1: - # drop right up corner - slash_scores = slash_scores[..., : -last_q_size + 1] - slash_scores[..., -100:] = torch.inf - - head_slash_size = heads_slash_size[head_i] - head_slash_size = min(head_slash_size, vertical.size(-1)) - slash_topk = torch.topk(slash_scores, head_slash_size, -1).indices - # （nheads, max_topk） - slash_topk_buffer[head_i, :head_slash_size] = slash_topk - - # reset heads topk - heads_slash_size[head_i] = head_slash_size - heads_vertical_size[head_i] = min( - heads_vertical_size[head_i], max_vertical_topk - ) - - # store - vertical_buffer = torch.full( - (n_heads, max_vertical_topk), - int32_max, - dtype=torch.int64, - device=q.device, - ) - slash_buffer = torch.full( - (n_heads, max_slash_topk), - int32_min, - dtype=torch.int64, - device=q.device, - ) - succ_vertical_buffer = torch.full( - (n_heads, max_vertical_topk), - int32_max, - dtype=torch.int64, - device=q.device, - ) - succ_slash_buffer = torch.full( - (n_heads, max_slash_topk), - int32_min, - dtype=torch.int64, - device=q.device, - ) - inter_vertical_buffer = torch.full( - (n_heads, max_vertical_topk), - int32_max, - dtype=torch.int64, - device=q.device, - ) - inter_slash_buffer = torch.full( - (n_heads, max_slash_topk), - int32_min, - dtype=torch.int64, - device=q.device, - ) - - vertical_size_buffer = torch.empty( - size=(n_heads,), dtype=torch.int32, device=q.device - ) - slash_sizes_buffer = torch.empty( - size=(n_heads,), dtype=torch.int32, device=q.device - ) - succ_vertical_size_buffer = torch.empty( - size=(n_heads,), dtype=torch.int32, device=q.device - ) - succ_slash_sizes_buffer = torch.empty( - size=(n_heads,), dtype=torch.int32, device=q.device - ) - inter_vertical_size_buffer = torch.empty( - size=(n_heads,), dtype=torch.int32, device=q.device - ) - inter_slash_sizes_buffer = torch.empty( - size=(n_heads,), dtype=torch.int32, device=q.device - ) - - for head_i in range(n_heads): - vertical_topk = vertical_topk_buffer[ - head_i, : heads_vertical_size[head_i] - ] - # intra - intra_vertical_indices = ( - vertical_topk[vertical_topk >= prev_chunk_end_pos] - - prev_chunk_end_pos - ) - if intra_vertical_indices.nelement() == 0: - intra_vertical_indices = _sparse_fallback_indices( - k_states_intra.size(0), - heads_vertical_size[head_i], - device=intra_vertical_indices.device, - ) - slash_topk = slash_topk_buffer[head_i, : heads_slash_size[head_i]] - intra_slash_indices = (qk.size(-1) - 1) - slash_topk[ - slash_topk >= prev_chunk_end_pos - ] - if intra_slash_indices.nelement() == 0: - intra_slash_indices = _sparse_fallback_indices( - k_states_intra.size(0), - heads_slash_size[head_i], - device=intra_vertical_indices.device, - ) - # fill buffer - v_count = intra_vertical_indices.nelement() - s_count = intra_slash_indices.nelement() - vertical_size_buffer[head_i] = v_count - slash_sizes_buffer[head_i] = s_count - vertical_buffer[head_i, :v_count].copy_(intra_vertical_indices) - slash_buffer[head_i, :s_count].copy_(intra_slash_indices) - # succ - if prev_chunk_end_pos - chunk_len >= 0: - succ_vertical_indices = vertical_topk[ - (vertical_topk < prev_chunk_end_pos) - & (vertical_topk >= prev_chunk_end_pos - chunk_len) - ] - (prev_chunk_end_pos - chunk_len) - # TODO: support no vertical - if succ_vertical_indices.nelement() == 0: - succ_vertical_indices = _sparse_fallback_indices( - k_states_succ.size(0), - heads_vertical_size[head_i], - device=intra_vertical_indices.device, - ) - succ_slash_indices = ( - prev_chunk_end_pos + (qend - qbegin) - 1 - ) - slash_topk[ - ( - (slash_topk >= (prev_chunk_end_pos - chunk_len)) - & (slash_topk < (prev_chunk_end_pos + (qend - qbegin))) - ) - ] - if succ_slash_indices.nelement() == 0: - succ_slash_indices = _sparse_fallback_indices( - k_states_succ.size(0), - heads_slash_size[head_i], - device=intra_vertical_indices.device, - ) - # fill buffer - v_count = succ_vertical_indices.nelement() - s_count = succ_slash_indices.nelement() - succ_vertical_size_buffer[head_i] = v_count - succ_slash_sizes_buffer[head_i] = s_count - succ_vertical_buffer[head_i, :v_count].copy_( - succ_vertical_indices - ) - succ_slash_buffer[head_i, :s_count].copy_(succ_slash_indices) - - if prev_chunk_end_pos - 2 * chunk_len >= 0: - inter_vertical_indices = vertical_topk[ - vertical_topk < prev_chunk_end_pos - chunk_len - ] - - if inter_vertical_indices.nelement() == 0: - inter_vertical_indices = _sparse_fallback_indices( - k_states_inter.size(0), - heads_vertical_size[head_i], - device=intra_vertical_indices.device, - ) - inter_slash_indices = ( - prev_chunk_end_pos - chunk_len + (qend - qbegin) - 1 - ) - slash_topk[ - slash_topk - < (prev_chunk_end_pos - chunk_len + (qend - qbegin)) - ] - if inter_slash_indices.nelement() == 0: - inter_slash_indices = _sparse_fallback_indices( - k_states_inter.size(0), - heads_slash_size[head_i], - device=intra_vertical_indices.device, - ) - # fill buffer - v_count = inter_vertical_indices.nelement() - s_count = inter_slash_indices.nelement() - inter_vertical_size_buffer[head_i] = v_count - inter_slash_sizes_buffer[head_i] = s_count - inter_vertical_buffer[head_i, :v_count].copy_( - inter_vertical_indices - ) - inter_slash_buffer[head_i, :s_count].copy_(inter_slash_indices) - else: - intra_vertical_indices, intra_slash_indices = None, None - succ_vertical_indices, succ_slash_indices = None, None - inter_vertical_indices, inter_slash_indices = None, None - - if sparse_attn_enabled: - flash_result = self._do_flash_attn( - q_states_intra, - k_states_intra, - v_states_intra, - softmax_scale=softmax_scale, - causal=True, - stage="intra", - vertical_indices=vertical_buffer, - slash_indices=slash_buffer, - vertical_indices_count=vertical_size_buffer, - slash_indices_count=slash_sizes_buffer, - mergehead_softmax_scale=softmax_scale, - sparse_attn_enabled=sparse_attn_enabled, - ) - else: - flash_result = self._do_flash_attn( - q_states_intra, - k_states_intra, - v_states_intra, - softmax_scale=softmax_scale, - causal=True, - stage="intra", - vertical_indices=intra_vertical_indices, - slash_indices=intra_slash_indices, - sparse_attn_enabled=sparse_attn_enabled, - ) - flash_per_chunk.append(flash_result) - - if prev_chunk_end_pos - chunk_len >= 0: - if sparse_attn_enabled: - flash_result = self._do_flash_attn( - q_states_succ, - k_states_succ, - v_states_succ, - softmax_scale=softmax_scale, - causal=False, - stage="succ", - vertical_indices=succ_vertical_buffer, - slash_indices=succ_slash_buffer, - vertical_indices_count=succ_vertical_size_buffer, - slash_indices_count=succ_slash_sizes_buffer, - mergehead_softmax_scale=softmax_scale, - sparse_attn_enabled=sparse_attn_enabled, - ) - else: - flash_result = self._do_flash_attn( - q_states_succ, - k_states_succ, - v_states_succ, - softmax_scale=softmax_scale, - causal=False, - stage="succ", - vertical_indices=succ_vertical_indices, - slash_indices=succ_slash_indices, - sparse_attn_enabled=sparse_attn_enabled, - ) - flash_per_chunk.append(flash_result) - - if prev_chunk_end_pos - chunk_len * 2 >= 0: - if sparse_attn_enabled: - flash_result = self._do_flash_attn( - q_states_inter, - k_states_inter, - v_states_inter, - softmax_scale=softmax_scale, - causal=False, - stage="inter", - vertical_indices=inter_vertical_buffer, - slash_indices=inter_slash_buffer, - vertical_indices_count=inter_vertical_size_buffer, - slash_indices_count=inter_slash_sizes_buffer, - mergehead_softmax_scale=softmax_scale, - sparse_attn_enabled=sparse_attn_enabled, - ) - else: - flash_result = self._do_flash_attn( - q_states_inter, - k_states_inter, - v_states_inter, - softmax_scale=softmax_scale, - causal=False, - stage="inter", - vertical_indices=inter_vertical_indices, - slash_indices=inter_slash_indices, - sparse_attn_enabled=sparse_attn_enabled, - ) - flash_per_chunk.append(flash_result) - - flash_results.append(flash_per_chunk) - begin = end - - attn_output = self._merge_attn_outputs(flash_results) - del flash_results - return attn_output - - def _do_flash_attn( - self, - query_states: torch.Tensor, - key_states: torch.Tensor, - value_states: torch.Tensor, - softmax_scale: float, - causal: bool = True, - max_seqlen_k: Optional[int] = None, - stage: str = "intra", - vertical_indices: Optional[torch.Tensor] = None, - slash_indices: Optional[torch.Tensor] = None, - vertical_indices_count: Optional[torch.Tensor] = None, - slash_indices_count: Optional[torch.Tensor] = None, - mergehead_softmax_scale: Optional[float] = None, - sparse_attn_enabled: Optional[bool] = False, - ): - if max_seqlen_k is None: - max_seqlen_k = key_states.shape[0] - - q_len = query_states.shape[0] - q_heads = query_states.shape[1] - h_dim = query_states.shape[-1] - - if sparse_attn_enabled: - assert slash_indices is not None - if stage == "intra": - assert causal - else: - assert not causal - - query_states = query_states.unsqueeze(0).transpose(1, 2) - key_states = key_states.unsqueeze(0).transpose(1, 2) - value_states = value_states.unsqueeze(0).transpose(1, 2) - - q = query_states - k = key_states - v = value_states - - if vertical_indices_count is not None and slash_indices_count is not None: - assert mergehead_softmax_scale is not None - - res, s_lse = _vertical_slash_sparse_attention( - q, - k, - v, - vertical_indices, - slash_indices, - mergehead_softmax_scale, - causal=causal, - stage=stage, - vertical_indices_count=vertical_indices_count, - slash_indices_count=slash_indices_count, - ) - res = res.view(q_heads, q_len, h_dim).transpose( - 0, 1 - ) # (qlen,nhead,h_dim) - s_lse = ( - s_lse.view(q_heads, q_len, 1).squeeze(-1).unsqueeze(0).float() - ) # (1, nhead,qlen) - else: - res, s_lse = _vertical_slash_sparse_attention( - q, - k, - v, - vertical_indices, - slash_indices, - softmax_scale, - causal=causal, - stage=stage, - ) - res = res.view(q_len, q_heads, h_dim) - s_lse = s_lse.view(q_len, q_heads, 1).transpose(0, 2).float() - return res, s_lse - - output, softmax_lse, *rest = flash_attn_varlen_func( - q=query_states, - k=key_states, - v=value_states, - softmax_scale=softmax_scale, - cu_seqlens_q=torch.tensor( - [0, query_states.shape[0]], - dtype=torch.int32, - device=query_states.device, - ), - max_seqlen_q=query_states.shape[0], - cu_seqlens_k=torch.tensor( - [0, max_seqlen_k], dtype=torch.int32, device=query_states.device - ), - max_seqlen_k=max_seqlen_k, - causal=causal, - return_softmax_lse=True, - ) - softmax_lse = softmax_lse.view(q_len, q_heads, 1).transpose(0, 2).float() - return output, softmax_lse - - def _merge_attn_outputs( - self, - flash_results: List[List[Tuple[torch.Tensor, torch.Tensor]]], - return_lse: Optional[bool] = False, - ) -> torch.Tensor: - attn_outputs_all = [] - logits_all = [] - - for flash_per_chunk in flash_results: - if len(flash_per_chunk) == 1: - attn_outputs_all.append(flash_per_chunk[0][0]) - if return_lse: - logits_all.append(flash_per_chunk[0][1]) - continue - - attn_outputs = torch.stack( - [flash_attn_output[0] for flash_attn_output in flash_per_chunk] - ) - logits = torch.stack( - [flash_attn_output[1] for flash_attn_output in flash_per_chunk] - ) - logits = logits.to(torch.float32) - - if return_lse: - max_val = torch.max(logits, dim=0).values - diff = torch.abs(logits[0] - logits[1]) - log_sum_exp = max_val + torch.log1p(torch.exp(-diff)) - logits_all.append(log_sum_exp) - - max_logits = torch.max(logits, dim=0).values - stable_logits = logits - max_logits.unsqueeze(0) - lse_s = torch.exp(stable_logits).detach() - lse_sum = torch.sum(lse_s, dim=0) - lse_s /= lse_sum - attn_outputs *= lse_s.unsqueeze(-1).transpose(2, 3).squeeze(1) - attn_outputs_all.append(attn_outputs.sum(dim=0)) - - if return_lse: - return (torch.cat(attn_outputs_all, dim=0), torch.cat(logits_all, dim=-1)) - else: - return torch.cat(attn_outputs_all, dim=0) - - def _dual_chunk_flash_attn_decoding( - self, - query: torch.Tensor, - query_succ: torch.Tensor, - query_inter: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - block_table: torch.Tensor, - cache_seqlens: torch.Tensor, - softmax_scale: float, - causal: bool, - chunk_size: int, - local_size: int, - original_max_position_embeddings: int, - decode_meta: DualChunkFlashAttentionMetadata, - ): - if not causal: - raise ValueError("Dual Chunk Attention does not support causal=False") - - block_size = value_cache.shape[1] - chunk_len = chunk_size - local_size - if chunk_len % block_size != 0: - raise ValueError("chunk_len must be divisible by block_size.") - if original_max_position_embeddings > 0: - assert decode_meta.scaling_factor is not None - scaling_factor = decode_meta.scaling_factor - query = (query * scaling_factor.view(-1, 1, 1, 1)).to( - query.dtype - ) # possible for numerical issue, need to fused in the kernel - query_succ = (query_succ * scaling_factor.view(-1, 1, 1, 1)).to(query.dtype) - query_inter = (query_inter * scaling_factor.view(-1, 1, 1, 1)).to( - query.dtype - ) - outputs_list = [] - softmax_lses_list = [] - - # intra-attention - intra_output, intra_softmax_lse = ( - self._dual_chunk_flash_attn_decoding_with_exp_sums( - query, - key_cache, - value_cache, - decode_meta.block_tables_intra, - decode_meta.seq_lens_intra, - softmax_scale, - causal=False, - ) - ) - outputs_list.append(intra_output) - softmax_lses_list.append(intra_softmax_lse) - - # succ-attention - if decode_meta.max_seq_len_succ: - succ_output, succ_softmax_lse = ( - self._dual_chunk_flash_attn_decoding_with_exp_sums( - query_succ, - key_cache, - value_cache, - decode_meta.block_tables_succ, - decode_meta.seq_lens_succ, - softmax_scale, - causal=False, - ) - ) - outputs_list.append(succ_output) - softmax_lses_list.append(succ_softmax_lse) - - # inter-attention - if decode_meta.max_seq_len_inter: - inter_output, inter_softmax_lse = ( - self._dual_chunk_flash_attn_decoding_with_exp_sums( - query_inter, - key_cache, - value_cache, - block_table, - decode_meta.seq_lens_inter, - softmax_scale, - causal=False, - ) - ) - outputs_list.append(inter_output) - softmax_lses_list.append(inter_softmax_lse) - outputs = torch.stack(outputs_list, dim=0) - del outputs_list - softmax_lses = torch.stack(softmax_lses_list, dim=0).to(torch.float32) - del softmax_lses_list - max_logits = torch.max(softmax_lses, dim=0).values - stable_logits = softmax_lses - max_logits.unsqueeze(0) - lse_s = torch.exp(stable_logits).detach() - lse_sum = torch.sum(lse_s, dim=0) - lse_s /= lse_sum - outputs *= lse_s.unsqueeze(-1).transpose(2, 3) - return outputs.sum(0) - - def _dual_chunk_flash_attn_decoding_with_exp_sums( - self, - query: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - block_table: torch.Tensor, - cache_seqlens: torch.Tensor, - softmax_scale: float, - causal: bool, - ): - out, softmax_lse, *rest_expand = flash_attn_with_kvcache( - q=query, - k_cache=key_cache, - v_cache=value_cache, - page_table=block_table, - cache_seqlens=cache_seqlens, - softmax_scale=softmax_scale, - causal=causal, - return_softmax_lse=True, - ) - mask = cache_seqlens == 0 - out[mask] = 0 - softmax_lse[mask] = -float("inf") - return out, softmax_lse - - -def _sparse_fallback_indices( - seq_len: int, max_count: int, device: torch.device -) -> torch.Tensor: - count = min(int(max_count), seq_len) - if count <= 0: - return torch.empty(0, dtype=torch.int64, device=device) - step = max(1, math.ceil(seq_len / count)) - return torch.arange(0, seq_len, step, dtype=torch.int64, device=device)[:count] - - -def _vertical_slash_sparse_attention( - query: torch.Tensor, # [BATCH, N_HEADS, N_CTX, D_HEAD] - key: torch.Tensor, # [BATCH, N_HEADS, N_KV_CTX, D_HEAD] - value: torch.Tensor, # [BATCH, N_HEADS, N_KV_CTX, D_HEAD] - v_idx: torch.Tensor, # [BATCH, N_HEADS, NNZ_V] - s_idx: torch.Tensor, # [BATCH, N_HEADS, NNZ_S] - softmax_scale: float, - causal: bool = True, - stage: str = "intra", - block_size_M: int = 64, - block_size_N: int = 64, - vertical_indices_count: torch.Tensor = None, # [N_HEADS,] - slash_indices_count: torch.Tensor = None, -): - if stage == "intra": - assert causal - else: - assert not causal - - batch_size, num_heads, context_size, head_dim = query.shape - _, _, kv_seq_len, _ = key.shape - - if head_dim not in [16, 32, 64, 128, 256, 512]: - target_dim = 2 ** math.ceil(math.log2(head_dim)) - head_dim - query = F.pad(query, [0, target_dim, 0, 0, 0, 0, 0, 0]) - key = F.pad(key, [0, target_dim, 0, 0, 0, 0, 0, 0]) - value = F.pad(value, [0, target_dim, 0, 0, 0, 0, 0, 0]) - - v_idx = ( - v_idx.to(torch.int32) - .reshape((batch_size, num_heads, -1)) - .sort(dim=-1, descending=False)[0] - ) - s_idx = ( - s_idx.to(torch.int32) - .reshape((batch_size, num_heads, -1)) - .sort(dim=-1, descending=True)[0] - ) - q_seqlens = torch.tensor([context_size], dtype=torch.int32, device=query.device) - kv_seqlens = torch.tensor([kv_seq_len], dtype=torch.int32, device=query.device) - - if vertical_indices_count is not None and slash_indices_count is not None: - ( - block_count, - block_offset, - column_count, - column_index, - ) = convert_vertical_slash_indexes_mergehead( - q_seqlens, - kv_seqlens, - v_idx, - s_idx, - vertical_indices_count, - slash_indices_count, - context_size, - block_size_M, - block_size_N, - causal, - ) - else: - ( - block_count, - block_offset, - column_count, - column_index, - ) = convert_vertical_slash_indexes( - q_seqlens, - kv_seqlens, - v_idx, - s_idx, - context_size, - block_size_M, - block_size_N, - causal, - ) - - q = query.transpose(1, 2).contiguous() - k = key.transpose(1, 2).contiguous() - v = value.transpose(1, 2).contiguous() - out, lse = sparse_attn_func( - q, - k, - v, - block_count, - block_offset, - column_count, - column_index, - causal=causal, - softmax_scale=softmax_scale, - return_softmax_lse=True, - ) - out = out.transpose(1, 2).contiguous() - softmax_lse = lse.reshape(*lse.shape, 1) - return (out[..., :context_size, :head_dim], softmax_lse[..., :context_size, :]) - - -def _sum_all_diagonal_matrix(mat: torch.tensor): - h, n, m = mat.shape - # Zero matrix used for padding - zero_mat = torch.zeros((h, n, n), device=mat.device) - # pads the matrix on left and right - mat_padded = torch.cat((zero_mat, mat, zero_mat), -1) - # Change the strides - mat_strided = mat_padded.as_strided( - (1, n, n + m), (n * (2 * n + m), 2 * n + m + 1, 1) - ) - # Sums the resulting matrix's columns - sum_diags = torch.sum(mat_strided, 1) - return sum_diags[:, 1:] # drop left bottom corner - - -def _get_block(block_table: torch.Tensor, block_size: int, begin: int, end: int): - begin_block = begin // block_size - end_block = (end - 1) // block_size + 1 - return block_table[begin_block:end_block] diff --git a/python/sglang/srt/layers/quantization/__init__.py b/python/sglang/srt/layers/quantization/__init__.py index 646165637..0f5829a6f 100644 --- a/python/sglang/srt/layers/quantization/__init__.py +++ b/python/sglang/srt/layers/quantization/__init__.py @@ -23,7 +23,6 @@ from sglang.srt.layers.quantization.gguf import GGUFConfig from sglang.srt.layers.quantization.gptq import ( CPUGPTQConfig, GPTQAscendConfig, - GPTQConfig, GPTQMarlinConfig, GPTQXPUConfig, ) @@ -75,7 +74,6 @@ BASE_QUANTIZATION_METHODS: Dict[str, Type[QuantizationConfig]] = { "awq_marlin": AWQMarlinConfig, "bitsandbytes": BitsAndBytesConfig, "gguf": GGUFConfig, - "gptq": GPTQConfig, "gptq_marlin": GPTQMarlinConfig, "moe_wna16": MoeWNA16Config, "compressed-tensors": CompressedTensorsConfig, @@ -116,6 +114,18 @@ if is_npu(): ) +if is_cpu(): + # Plain GPTQ is CUDA-only in name: the kernel is gone, but the Intel AMX + # path below is untouched. `get_quantization_config` rejects anything + # missing from this registry before it can consult CPU_QUANTIZATION_METHODS, + # so the key has to exist here for the CPU path to stay reachable. + BASE_QUANTIZATION_METHODS.update( + { + "gptq": CPUGPTQConfig, + } + ) + + if is_xpu(): BASE_QUANTIZATION_METHODS.update( { diff --git a/python/sglang/srt/layers/quantization/gptq/gptq.py b/python/sglang/srt/layers/quantization/gptq/gptq.py index 123789bad..d53f62844 100644 --- a/python/sglang/srt/layers/quantization/gptq/gptq.py +++ b/python/sglang/srt/layers/quantization/gptq/gptq.py @@ -661,18 +661,8 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase): # Register fake implementations for torch.compile support. The decorator is a # no-op when the custom op is unavailable on the current platform. -@register_fake_if_exists("sgl_kernel::gptq_gemm") -def _(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, use_shuffle, bit): - return a.new_empty((a.shape[0], b_q_weight.shape[-1]), dtype=a.dtype) - - @register_fake_if_exists("sgl_kernel::gptq_marlin_repack") def _(b_q_weight, perm, size_k, size_n, num_bits): return b_q_weight.new_empty( (size_k // 16, size_n * (num_bits // 2)), dtype=b_q_weight.dtype ) - - -@register_fake_if_exists("sgl_kernel::gptq_shuffle") -def _(q_weight, q_perm, bit): - return diff --git a/python/sglang/srt/layers/quantization/gptq/schemes/gptq_linear.py b/python/sglang/srt/layers/quantization/gptq/schemes/gptq_linear.py index 8f83426c0..45ba29a34 100644 --- a/python/sglang/srt/layers/quantization/gptq/schemes/gptq_linear.py +++ b/python/sglang/srt/layers/quantization/gptq/schemes/gptq_linear.py @@ -29,12 +29,11 @@ class GPTQLinearScheme(GPTQLinearSchemeBase): self.kernel = self._init_kernel(quant_config) def _init_kernel(self, quant_config: GPTQConfig): - from sglang.srt.hardware_backend.gpu.quantization.gptq_kernels import ( - GPTQLinearKernel, + raise RuntimeError( + "The non-Marlin GPTQ CUDA kernel has been removed. Use " + "quantization='gptq_marlin' (or a Marlin-compatible checkpoint) instead." ) - return GPTQLinearKernel(quant_config) - def create_weights( self, layer: torch.nn.Module, diff --git a/python/sglang/srt/models/bailing_moe_linear.py b/python/sglang/srt/models/bailing_moe_linear.py index 07da00d68..43d880b76 100644 --- a/python/sglang/srt/models/bailing_moe_linear.py +++ b/python/sglang/srt/models/bailing_moe_linear.py @@ -98,7 +98,7 @@ if _use_aiter_gfx95: pass if _is_cuda: - from sgl_kernel import awq_dequantize + from sglang.kernels.ops.quantization.awq_dequantize import awq_dequantize elif _is_cpu and _is_cpu_amx_available: pass elif _is_hip: diff --git a/python/sglang/srt/models/bailing_moe_v3.py b/python/sglang/srt/models/bailing_moe_v3.py index a73c2e0d5..bac6a39d4 100644 --- a/python/sglang/srt/models/bailing_moe_v3.py +++ b/python/sglang/srt/models/bailing_moe_v3.py @@ -99,7 +99,7 @@ from sglang.srt.utils import ( _is_fp8_fnuz = is_fp8_fnuz() if _is_cuda: - from sgl_kernel import awq_dequantize + from sglang.kernels.ops.quantization.awq_dequantize import awq_dequantize elif _is_cpu and _is_cpu_amx_available: pass elif _is_hip: diff --git a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py index 8543c3fdf..a730a8013 100644 --- a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py +++ b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py @@ -156,10 +156,6 @@ def handle_attention_flashmla(attn, forward_batch): return _handle_attention_backend(attn, forward_batch, "flashmla") -def handle_attention_cutlass_mla(attn, forward_batch): - return _handle_attention_backend(attn, forward_batch, "cutlass_mla") - - def handle_attention_fa4(attn, forward_batch): # FA4 absorbed MLA feeds q_nope through the qv argument, which # flash_attn.cute only implements on SM100/SM110 (not SM120); keep the @@ -268,7 +264,6 @@ AttentionBackendRegistry.register("ascend", handle_attention_ascend) AttentionBackendRegistry.register("flashinfer", handle_attention_flashinfer) AttentionBackendRegistry.register("fa3", handle_attention_fa3) AttentionBackendRegistry.register("flashmla", handle_attention_flashmla) -AttentionBackendRegistry.register("cutlass_mla", handle_attention_cutlass_mla) AttentionBackendRegistry.register("fa4", handle_attention_fa4) AttentionBackendRegistry.register("trtllm_mla", handle_attention_trtllm_mla) AttentionBackendRegistry.register("tokenspeed_mla", handle_attention_tokenspeed_mla) diff --git a/python/sglang/srt/models/deepseek_common/utils.py b/python/sglang/srt/models/deepseek_common/utils.py index 45064be29..ceefce6d8 100644 --- a/python/sglang/srt/models/deepseek_common/utils.py +++ b/python/sglang/srt/models/deepseek_common/utils.py @@ -68,7 +68,6 @@ FORWARD_ABSORB_CORE_ATTENTION_BACKENDS = [ "dsa", "nsa", # Deprecated alias for "dsa" "flashinfer", - "cutlass_mla", "trtllm_mla", "cutedsl_mla", "tokenspeed_mla", @@ -103,9 +102,10 @@ def awq_dequantize_func(): - None if the current device is not supported. """ if _is_cuda: - from sgl_kernel import awq_dequantize + from sglang.kernel_api_logging import debug_kernel_api + from sglang.kernels.ops.quantization.awq_dequantize import awq_dequantize - return awq_dequantize + return debug_kernel_api(awq_dequantize, op_name="DeepseekCommon.awq_dequantize") elif _is_hip: from sglang.kernels.kernel_api_logging import debug_kernel_api from sglang.kernels.ops.quantization.awq_triton import ( diff --git a/python/sglang/srt/models/longcat_flash.py b/python/sglang/srt/models/longcat_flash.py index e2571ba28..711094135 100644 --- a/python/sglang/srt/models/longcat_flash.py +++ b/python/sglang/srt/models/longcat_flash.py @@ -116,7 +116,7 @@ _is_cpu = is_cpu() _device_sm = get_device_sm() if _is_cuda: - from sgl_kernel import awq_dequantize + from sglang.kernels.ops.quantization.awq_dequantize import awq_dequantize elif _is_cpu and _is_cpu_amx_available: pass elif _is_hip: diff --git a/python/sglang/srt/models/longcat_flash_nextn.py b/python/sglang/srt/models/longcat_flash_nextn.py index 0a799b56d..40c1bc909 100644 --- a/python/sglang/srt/models/longcat_flash_nextn.py +++ b/python/sglang/srt/models/longcat_flash_nextn.py @@ -92,7 +92,7 @@ _is_cpu = is_cpu() _device_sm = get_device_sm() if _is_cuda: - from sgl_kernel import awq_dequantize + from sglang.kernels.ops.quantization.awq_dequantize import awq_dequantize elif _is_cpu and _is_cpu_amx_available: pass elif _is_hip: diff --git a/python/sglang/srt/models/sarvam_moe.py b/python/sglang/srt/models/sarvam_moe.py index 63373a9a2..90d877c9b 100644 --- a/python/sglang/srt/models/sarvam_moe.py +++ b/python/sglang/srt/models/sarvam_moe.py @@ -115,7 +115,7 @@ class AttnForwardMethod(IntEnum): SEPARATE_ROPE_BACKENDS = frozenset( - ["fa3", "flashinfer", "dsa", "nsa", "cutlass_mla", "trtllm_mla"] + ["fa3", "flashinfer", "dsa", "nsa", "trtllm_mla"] # "nsa" is a deprecated alias for "dsa" ) CONCAT_ROPE_BACKENDS = frozenset(["flashmla", "triton"]) diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py index a69346934..77bb8e5d4 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py @@ -1190,8 +1190,8 @@ def dsa_impl_capability(impl: str) -> tuple[bool, str]: # TRT-LLM Gen FMHA / MLA require Blackwell SM10.0 (B200 NVL). # SM10.3 (GB300) raises "Missing TRTLLM-GEN kernel" at runtime because # the kernel binary in the container isn't compiled for sm_103. - # Require exactly SM10.0 (same constraint as cutlass_mla) until the - # container ships sm_103-compiled TRTLLM-GEN kernels. + # Require exactly SM10.0 until the container ships sm_103-compiled + # TRTLLM-GEN kernels. if major != 10 or minor != 0: return ( False, diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py deleted file mode 100644 index 3a62e1031..000000000 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py +++ /dev/null @@ -1,1605 +0,0 @@ -from dataclasses import dataclass -from types import SimpleNamespace - -import torch -from torch import nn - -from sglang.srt.distributed.parallel_state_wrapper import ParallelState -from sglang.srt.layers.attention import ( - dual_chunk_flashattention_backend as _dual_chunk_backend, -) -from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS -from sglang.srt.layers.radix_attention import RadixAttention -from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, ReqToTokenPool -from sglang.srt.model_executor.cuda_graph_config import ( - Backend, - CudaGraphConfig, - PhaseConfig, -) -from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode -from sglang.srt.model_executor.forward_context import ForwardContext, forward_context -from sglang.srt.model_executor.model_runner import ModelRunner -from sglang.srt.runtime_context import get_context, get_parallel - -from .dense_attention import ( - DEFAULT_DEVICE, - DEFAULT_DTYPE, - DEFAULT_HEAD_DIM, - DEFAULT_HIDDEN_SIZE, - DEFAULT_MAX_CONTEXT_LEN, - DENSE_ATOL, - DENSE_RTOL, - DenseAttentionCase, - ReferenceDenseAttention, - _copy_dense_weights, - _expand_gqa, - _make_forward_batch, - _populate_prefix_kv, - _split_by_lens, -) - -DUAL_CHUNK_CONFIG = { - "chunk_size": 64, - "local_size": 16, - "original_max_position_embeddings": 32768, - "sparse_attention_enabled": False, -} -DUAL_CHUNK_SPARSE_ALL_COLUMN_CONFIG = { - **DUAL_CHUNK_CONFIG, - "sparse_attention_enabled": True, - "sparse_attention_threshold": 0, - "sparse_attention_last_q": 16, - "sparse_attention_config": { - 0: {str(head_id): ("vertical_and_slash", 16, 16, None) for head_id in range(4)} - }, -} -# Same vertical/slash sizes as all-column, but with a threshold so short -# sequences bypass the sparse kernel and fall back to dense prefill. This -# exercises the `current_orig_seq_len > self.sparse_attention_threshold` gate. -DUAL_CHUNK_SPARSE_THRESHOLD_GATED_CONFIG = { - **DUAL_CHUNK_SPARSE_ALL_COLUMN_CONFIG, - "sparse_attention_threshold": 100, -} -# Sub-context-window sparse: vertical_size + slash_size < intra K count, so -# the kernel's vertical+slash topk genuinely prunes (the union of selected -# columns + slashes does NOT cover every K column). The selection is -# deterministic but content-aware; this case compares the sgl-kernel sparse -# attention output against a torch implementation that consumes the same -# block/column metadata. -DUAL_CHUNK_SPARSE_SUB_WINDOW_CONFIG = { - **DUAL_CHUNK_CONFIG, - "sparse_attention_enabled": True, - "sparse_attention_threshold": 0, - "sparse_attention_last_q": 8, - "sparse_attention_config": { - 0: {str(head_id): ("vertical_and_slash", 4, 4, None) for head_id in range(4)} - }, -} - -# Unit tests run without distributed initialization. Sparse dual-chunk config -# lookup should see the single-rank default. -_parallel_override = get_parallel().override(tp_rank=0) -_parallel_override.__enter__() - - -@dataclass(frozen=True) -class DualChunkAttentionCase(DenseAttentionCase): - pass - - -def make_dual_chunk_cases(backend: str) -> tuple[DualChunkAttentionCase, ...]: - common = dict(backend=backend, num_heads=4, num_kv_heads=4) - return ( - DualChunkAttentionCase( - name="dual_chunk_extend_page_size_1", - forward_mode=ForwardMode.EXTEND, - page_size=1, - prefix_lens=(2, 4), - extend_lens=(3, 1), - **common, - ), - DualChunkAttentionCase( - name="dual_chunk_extend_zero_prefix_exact_page", - forward_mode=ForwardMode.EXTEND, - page_size=16, - prefix_lens=(0,), - extend_lens=(16,), - **common, - ), - DualChunkAttentionCase( - name="dual_chunk_extend_cross_page_boundary", - forward_mode=ForwardMode.EXTEND, - page_size=16, - prefix_lens=(15,), - extend_lens=(2,), - **common, - ), - DualChunkAttentionCase( - name="dual_chunk_extend_ragged_page_boundary", - forward_mode=ForwardMode.EXTEND, - page_size=16, - prefix_lens=(0, 8, 16), - extend_lens=(15, 8, 1), - **common, - ), - DualChunkAttentionCase( - name="dual_chunk_decode_page_boundary", - forward_mode=ForwardMode.DECODE, - page_size=16, - prefix_lens=(14, 15, 16), - **common, - ), - DualChunkAttentionCase( - name="dual_chunk_extend_succ_chunk", - forward_mode=ForwardMode.EXTEND, - page_size=16, - prefix_lens=(46,), - extend_lens=(4,), - **common, - ), - DualChunkAttentionCase( - name="dual_chunk_decode_succ_chunk", - forward_mode=ForwardMode.DECODE, - page_size=16, - prefix_lens=(48,), - **common, - ), - DualChunkAttentionCase( - name="dual_chunk_extend_inter_chunk", - forward_mode=ForwardMode.EXTEND, - page_size=16, - prefix_lens=(94,), - extend_lens=(4,), - **common, - ), - DualChunkAttentionCase( - name="dual_chunk_decode_inter_chunk", - forward_mode=ForwardMode.DECODE, - page_size=16, - prefix_lens=(96,), - **common, - ), - DualChunkAttentionCase( - name="dual_chunk_gqa_decode_page_boundary", - forward_mode=ForwardMode.DECODE, - num_heads=4, - num_kv_heads=2, - page_size=16, - prefix_lens=(14, 15, 16), - backend=backend, - ), - DualChunkAttentionCase( - name="dual_chunk_gqa_decode_inter_chunk", - forward_mode=ForwardMode.DECODE, - num_heads=4, - num_kv_heads=2, - page_size=16, - prefix_lens=(96,), - backend=backend, - ), - ) - - -def make_dual_chunk_sparse_cases(backend: str) -> tuple[DualChunkAttentionCase, ...]: - common = dict(backend=backend, num_heads=4, num_kv_heads=4) - return ( - DualChunkAttentionCase( - name="dual_chunk_sparse_prefill_all_columns", - forward_mode=ForwardMode.EXTEND, - page_size=16, - prefix_lens=(0,), - extend_lens=(16,), - **common, - ), - # Multi-request batch within the first chunk: every request's seq_len <= 16 - # so the sparse path's per-request all-column selection still covers all keys - # and matches the dense reference. Exercises per-request `cu_seqlens_*` slicing - # in `_dual_chunk_flash_attn_prefill_func` under sparse enabled. - DualChunkAttentionCase( - name="dual_chunk_sparse_prefill_multi_request_first_chunk", - forward_mode=ForwardMode.EXTEND, - page_size=16, - prefix_lens=(0, 0), - extend_lens=(8, 12), - **common, - ), - # Page-boundary extend (prefix + extend crosses page=16) while staying within - # one chunk (chunk_size=64). Sparse path still sees <= 16 keys per request so - # last_q + vertical/slash select all → dense-equivalent. - DualChunkAttentionCase( - name="dual_chunk_sparse_prefill_cross_page_first_chunk", - forward_mode=ForwardMode.EXTEND, - page_size=16, - prefix_lens=(15,), - extend_lens=(1,), - **common, - ), - ) - - -def make_dual_chunk_sparse_threshold_gated_cases( - backend: str, -) -> tuple[DualChunkAttentionCase, ...]: - common = dict(backend=backend, num_heads=4, num_kv_heads=4) - return ( - # sparse_attention_enabled=True with threshold=100; with seq_len=16 the - # backend's `current_orig_seq_len > threshold` check should disable sparse - # per request and fall back to the dense chunk-flash kernel. The output - # must match the dense reference exactly. - DualChunkAttentionCase( - name="dual_chunk_sparse_threshold_gated_short_seq", - forward_mode=ForwardMode.EXTEND, - page_size=16, - prefix_lens=(0,), - extend_lens=(16,), - **common, - ), - ) - - -def make_dual_chunk_sparse_sub_window_cases( - backend: str, -) -> tuple[DualChunkAttentionCase, ...]: - common = dict(backend=backend, num_heads=4, num_kv_heads=4) - return ( - DualChunkAttentionCase( - name="dual_chunk_sparse_prefill_sub_window_seq128", - forward_mode=ForwardMode.EXTEND, - page_size=16, - prefix_lens=(0,), - extend_lens=(128,), - **common, - ), - ) - - -class TinyDualChunkModelConfig: - def __init__( - self, - *, - num_heads: int, - num_kv_heads: int, - head_dim: int, - hidden_size: int, - context_len: int, - dual_chunk_attention_config: dict | None = None, - ): - self.context_len = context_len - self.hidden_size = hidden_size - self.num_attention_heads = num_heads - self.num_key_value_heads = num_kv_heads - self.head_dim = head_dim - self.v_head_dim = head_dim - self.is_encoder_decoder = False - self.is_multimodal = False - self.is_generation = True - self.quantization = None - self.is_hybrid_swa = False - self.attention_chunk_size = None - self.sliding_window_size = None - self.hf_config = SimpleNamespace( - architectures=["TinyDualChunkForCausalLM"], - hidden_size=hidden_size, - num_attention_heads=num_heads, - num_key_value_heads=num_kv_heads, - head_dim=head_dim, - dual_chunk_attention_config=( - dual_chunk_attention_config or DUAL_CHUNK_CONFIG - ), - ) - self.hf_config.get_text_config = lambda: self.hf_config - self.hf_text_config = self.hf_config - self.linear_attn_registry_result = None - - def get_num_attention_heads(self, tp_size: int) -> int: - assert self.num_attention_heads % tp_size == 0 - return self.num_attention_heads // tp_size - - def get_max_num_attention_heads(self) -> int: - return self.num_attention_heads - - def get_num_kv_heads(self, tp_size: int, dcp_size: int = 1) -> int: - kv_tp_size = tp_size // dcp_size - assert self.num_key_value_heads % kv_tp_size == 0 - return self.num_key_value_heads // kv_tp_size - - -class DualChunkMockModelRunner(ModelRunner): - def __init__( - self, - *, - case: DualChunkAttentionCase, - model_config: TinyDualChunkModelConfig, - dtype: torch.dtype, - device: str, - max_context_len: int, - head_dim: int, - disable_cuda_graph: bool = True, - disable_piecewise_cuda_graph: bool = True, - runner_batch_size: int | None = None, - ): - pool_batch_size = runner_batch_size or case.batch_size - self.device = device - self.dtype = dtype - self.kv_cache_dtype = dtype - self.kv_cache_dtype_str = "auto" - # This runner's own resolved backends (production stamps these in - # ModelRunner.initialize); a draft runner would carry its own. - self.prefill_attention_backend_str = case.backend - self.decode_attention_backend_str = case.backend - self.draft_attention_backend = None - self.gpu_id = 0 - self.canary_manager = None - self.page_size = case.page_size - self.model_config = model_config - self.tp_size = 1 - self._kernel_warmed_up = True - self.dp_size = 1 - self.pp_size = 1 - self.ps = ParallelState.trivial() - self._server_args_override = get_context().override_server_args( - attention_backend=case.backend, - chunked_prefill_size=-1, - cuda_graph_config=CudaGraphConfig( - decode=PhaseConfig( - backend=Backend.DISABLED if disable_cuda_graph else Backend.FULL, - ), - prefill=PhaseConfig( - backend=( - Backend.DISABLED - if (disable_cuda_graph or disable_piecewise_cuda_graph) - else Backend.TC_PIECEWISE - ), - ), - ), - disable_radix_cache=False, - dp_size=1, - enable_dp_attention=False, - kv_cache_dtype="auto", - speculative_algorithm=None, - speculative_eagle_topk=0, - speculative_num_draft_tokens=0, - speculative_num_steps=0, - tp_size=1, - triton_attention_num_kv_splits=8, - triton_attention_split_tile_size=None, - ) - self.server_args = self._server_args_override.install() - self.req_to_token_pool = ReqToTokenPool( - size=pool_batch_size, - max_context_len=max_context_len, - device=device, - enable_memory_saver=False, - ) - max_token_loc = case.page_size + pool_batch_size * max_context_len - self.token_to_kv_pool = MHATokenToKVPool( - size=max_token_loc + case.page_size, - page_size=case.page_size, - dtype=dtype, - head_num=case.num_kv_heads, - head_dim=head_dim, - layer_num=1, - device=device, - enable_memory_saver=False, - enable_alt_stream=False, - ) - self.token_to_kv_pool_allocator = SimpleNamespace(page_size=case.page_size) - self.init_kv_index_translator() - self.attn_cp_size = 1 - self.attention_chunk_size = None - self.hisparse_coordinator = None - self.init_new_workspace = False - self.is_hybrid_swa = False - self.use_mla_backend = False - - @property - def hybrid_gdn_config(self): - return None - - @property - def hybrid_lightning_config(self): - return None - - @property - def kimi_linear_config(self): - return None - - @property - def linear_attn_model_spec(self): - return None - - @property - def mamba2_config(self): - return None - - @property - def mambaish_config(self): - return None - - -class ProjectedDualChunkAttention(nn.Module): - def __init__( - self, - *, - hidden_size: int, - num_heads: int, - num_kv_heads: int, - head_dim: int, - dtype: torch.dtype, - device: str, - ): - super().__init__() - self.hidden_size = hidden_size - self.num_heads = num_heads - self.num_kv_heads = num_kv_heads - self.head_dim = head_dim - self.q_proj = nn.Linear( - hidden_size, - num_heads * head_dim, - bias=False, - dtype=dtype, - device=device, - ) - self.q_succ_proj = nn.Linear( - hidden_size, - num_heads * head_dim, - bias=False, - dtype=dtype, - device=device, - ) - self.q_inter_proj = nn.Linear( - hidden_size, - num_heads * head_dim, - bias=False, - dtype=dtype, - device=device, - ) - self.q_succ_critical_proj = nn.Linear( - hidden_size, - num_heads * head_dim, - bias=False, - dtype=dtype, - device=device, - ) - self.q_inter_critical_proj = nn.Linear( - hidden_size, - num_heads * head_dim, - bias=False, - dtype=dtype, - device=device, - ) - self.k_proj = nn.Linear( - hidden_size, - num_kv_heads * head_dim, - bias=False, - dtype=dtype, - device=device, - ) - self.v_proj = nn.Linear( - hidden_size, - num_kv_heads * head_dim, - bias=False, - dtype=dtype, - device=device, - ) - self.o_proj = nn.Linear( - num_heads * head_dim, - hidden_size, - bias=False, - dtype=dtype, - device=device, - ) - self.attn = RadixAttention( - num_heads=num_heads, - head_dim=head_dim, - scaling=head_dim**-0.5, - num_kv_heads=num_kv_heads, - layer_id=0, - ) - - def project_qkv(self, hidden_states: torch.Tensor): - q = self.q_proj(hidden_states) - k = self.k_proj(hidden_states) - v = self.v_proj(hidden_states) - return q, k, v - - def project_dual_qkv(self, hidden_states: torch.Tensor): - q = self.q_proj(hidden_states) - q_succ = self.q_succ_proj(hidden_states) - q_inter = self.q_inter_proj(hidden_states) - q_succ_critical = self.q_succ_critical_proj(hidden_states) - q_inter_critical = self.q_inter_critical_proj(hidden_states) - k = self.k_proj(hidden_states) - v = self.v_proj(hidden_states) - return q, q_succ, q_inter, q_succ_critical, q_inter_critical, k, v - - def forward(self, hidden_states: torch.Tensor, forward_batch: ForwardBatch): - q, q_succ, q_inter, q_succ_critical, q_inter_critical, k, v = ( - self.project_dual_qkv(hidden_states) - ) - packed_q = torch.cat( - (q, q_succ, q_inter, q_succ_critical, q_inter_critical), dim=-1 - ) - attn_output = self.attn(packed_q, k, v, forward_batch) - return self.o_proj(attn_output) - - -class ReferenceDualChunkAttention(ReferenceDenseAttention): - def __init__( - self, - *, - hidden_size: int, - num_heads: int, - num_kv_heads: int, - head_dim: int, - dtype: torch.dtype, - device: str, - ): - super().__init__( - hidden_size=hidden_size, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - head_dim=head_dim, - dtype=dtype, - device=device, - ) - self.q_succ_proj = nn.Linear( - hidden_size, - num_heads * head_dim, - bias=False, - dtype=dtype, - device=device, - ) - self.q_inter_proj = nn.Linear( - hidden_size, - num_heads * head_dim, - bias=False, - dtype=dtype, - device=device, - ) - self.q_succ_critical_proj = nn.Linear( - hidden_size, - num_heads * head_dim, - bias=False, - dtype=dtype, - device=device, - ) - self.q_inter_critical_proj = nn.Linear( - hidden_size, - num_heads * head_dim, - bias=False, - dtype=dtype, - device=device, - ) - - def project_dual_qkv(self, hidden_states: torch.Tensor): - q = self.q_proj(hidden_states) - q_succ = self.q_succ_proj(hidden_states) - q_inter = self.q_inter_proj(hidden_states) - q_succ_critical = self.q_succ_critical_proj(hidden_states) - q_inter_critical = self.q_inter_critical_proj(hidden_states) - k = self.k_proj(hidden_states) - v = self.v_proj(hidden_states) - return q, q_succ, q_inter, q_succ_critical, q_inter_critical, k, v - - -@dataclass -class DualChunkAttentionFixture: - case: DualChunkAttentionCase - runner: DualChunkMockModelRunner - backend: object - actual_module: ProjectedDualChunkAttention - reference_module: ReferenceDualChunkAttention - forward_batch: ForwardBatch - prefix_hidden: list[torch.Tensor] - input_hidden: torch.Tensor - - -@dataclass(frozen=True) -class _DualChunkSparseStageSelection: - stage: str - q_len: int - kv_len: int - vertical_indices: tuple[tuple[int, ...], ...] - slash_indices: tuple[tuple[int, ...], ...] - - -def _set_orig_seq_lens(batch: ForwardBatch, case: DualChunkAttentionCase) -> None: - batch.orig_seq_lens = torch.tensor( - case.seq_lens, - dtype=torch.int32, - device=batch.seq_lens.device, - ) - - -def build_dual_chunk_attention_fixture( - testcase, - case: DualChunkAttentionCase, - *, - head_dim: int = DEFAULT_HEAD_DIM, - hidden_size: int = DEFAULT_HIDDEN_SIZE, - max_context_len: int = DEFAULT_MAX_CONTEXT_LEN, - dtype: torch.dtype = DEFAULT_DTYPE, - device: str = DEFAULT_DEVICE, - dual_chunk_attention_config: dict | None = None, - disable_cuda_graph: bool = True, - disable_piecewise_cuda_graph: bool = True, - runner_batch_size: int | None = None, - loc_layout: str = "shuffled_pages", -) -> DualChunkAttentionFixture: - max_context_len = max(max_context_len, max(case.seq_lens)) - if max_context_len % case.page_size: - max_context_len = ( - (max_context_len + case.page_size - 1) // case.page_size - ) * case.page_size - - seed = 3026 + len(case.name) + case.num_kv_heads - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - - model_config = TinyDualChunkModelConfig( - num_heads=case.num_heads, - num_kv_heads=case.num_kv_heads, - head_dim=head_dim, - hidden_size=hidden_size, - context_len=max_context_len, - dual_chunk_attention_config=dual_chunk_attention_config, - ) - runner = DualChunkMockModelRunner( - case=case, - model_config=model_config, - dtype=dtype, - device=device, - max_context_len=max_context_len, - head_dim=head_dim, - disable_cuda_graph=disable_cuda_graph, - disable_piecewise_cuda_graph=disable_piecewise_cuda_graph, - runner_batch_size=runner_batch_size, - ) - try: - backend = ATTENTION_BACKENDS[case.backend](runner) - except (AssertionError, ImportError, ModuleNotFoundError) as exc: - testcase.skipTest(f"{case.backend} backend is not available: {exc}") - - actual_module = ProjectedDualChunkAttention( - hidden_size=hidden_size, - num_heads=case.num_heads, - num_kv_heads=case.num_kv_heads, - head_dim=head_dim, - dtype=dtype, - device=device, - ) - reference_module = ReferenceDualChunkAttention( - hidden_size=hidden_size, - num_heads=case.num_heads, - num_kv_heads=case.num_kv_heads, - head_dim=head_dim, - dtype=dtype, - device=device, - ) - _copy_dual_chunk_weights(actual_module, reference_module) - prefix_hidden = [ - torch.randn(length, hidden_size, dtype=dtype, device=device) - for length in case.prefix_lens - ] - input_hidden = torch.randn( - case.num_input_tokens, - hidden_size, - dtype=dtype, - device=device, - ) - from .dense_attention import make_loc_fn as _dense_make_loc_fn - - loc_fn = _dense_make_loc_fn( - loc_layout, - batch_size=case.batch_size, - seq_lens=case.seq_lens, - prefix_lens=case.prefix_lens, - page_size=case.page_size, - max_context_len=max_context_len, - seed=seed, - ) - forward_batch = _make_forward_batch( - case, - runner, - max_context_len=max_context_len, - device=device, - loc_fn=loc_fn, - ) - _set_orig_seq_lens(forward_batch, case) - _populate_prefix_kv( - actual_module, - case, - runner, - prefix_hidden, - max_context_len=max_context_len, - loc_fn=loc_fn, - ) - - return DualChunkAttentionFixture( - case=case, - runner=runner, - backend=backend, - actual_module=actual_module, - reference_module=reference_module, - forward_batch=forward_batch, - prefix_hidden=prefix_hidden, - input_hidden=input_hidden, - ) - - -def run_dual_chunk_fixture_eager(fixture: DualChunkAttentionFixture) -> torch.Tensor: - with torch.no_grad(), forward_context(ForwardContext(attn_backend=fixture.backend)): - fixture.backend.init_forward_metadata(fixture.forward_batch) - return fixture.actual_module(fixture.input_hidden, fixture.forward_batch) - - -def expected_dual_chunk_fixture_output( - fixture: DualChunkAttentionFixture, -) -> torch.Tensor: - return _dual_chunk_attention_reference( - fixture.reference_module, - fixture.case, - fixture.prefix_hidden, - fixture.input_hidden, - ) - - -def _copy_dual_chunk_weights( - actual: ProjectedDualChunkAttention, - reference: ReferenceDualChunkAttention, -) -> None: - _copy_dense_weights(actual, reference) - with torch.no_grad(): - reference.q_succ_proj.weight.copy_(actual.q_succ_proj.weight) - reference.q_inter_proj.weight.copy_(actual.q_inter_proj.weight) - reference.q_succ_critical_proj.weight.copy_(actual.q_succ_critical_proj.weight) - reference.q_inter_critical_proj.weight.copy_( - actual.q_inter_critical_proj.weight - ) - - -def _dual_chunk_attention_reference( - module: ReferenceDualChunkAttention, - case: DualChunkAttentionCase, - prefix_hidden: list[torch.Tensor], - input_hidden: torch.Tensor, -) -> torch.Tensor: - dtype = input_hidden.dtype - q, q_succ, q_inter, _, _, k, v = module.project_dual_qkv(input_hidden) - q_parts = _split_by_lens( - q.view(-1, case.num_heads, module.head_dim), case.input_lens - ) - q_succ_parts = _split_by_lens( - q_succ.view(-1, case.num_heads, module.head_dim), case.input_lens - ) - q_inter_parts = _split_by_lens( - q_inter.view(-1, case.num_heads, module.head_dim), case.input_lens - ) - k_parts = _split_by_lens( - k.view(-1, case.num_kv_heads, module.head_dim), case.input_lens - ) - v_parts = _split_by_lens( - v.view(-1, case.num_kv_heads, module.head_dim), case.input_lens - ) - outputs = [] - chunk_len = DUAL_CHUNK_CONFIG["chunk_size"] - DUAL_CHUNK_CONFIG["local_size"] - - for req_idx, prefix in enumerate(prefix_hidden): - _, _, _, _, _, prefix_k, prefix_v = module.project_dual_qkv(prefix) - prefix_k = prefix_k.view(-1, case.num_kv_heads, module.head_dim) - prefix_v = prefix_v.view(-1, case.num_kv_heads, module.head_dim) - req_k = torch.cat([prefix_k, k_parts[req_idx]], dim=0) - req_v = torch.cat([prefix_v, v_parts[req_idx]], dim=0) - - for offset, query in enumerate(q_parts[req_idx]): - query_pos = case.prefix_lens[req_idx] + offset - current_chunk_start = (query_pos // chunk_len) * chunk_len - previous_chunk_start = current_chunk_start - chunk_len - groups = [ - ( - query, - req_k[current_chunk_start : query_pos + 1], - req_v[current_chunk_start : query_pos + 1], - ) - ] - - if previous_chunk_start >= 0: - groups.append( - ( - q_succ_parts[req_idx][offset], - req_k[previous_chunk_start:current_chunk_start], - req_v[previous_chunk_start:current_chunk_start], - ) - ) - - if previous_chunk_start > 0: - groups.append( - ( - q_inter_parts[req_idx][offset], - req_k[:previous_chunk_start], - req_v[:previous_chunk_start], - ) - ) - - score_parts = [] - value_parts = [] - for group_query, group_k, group_v in groups: - keys = _expand_gqa(group_k.movedim(0, 1), case.num_heads) - values = _expand_gqa(group_v.movedim(0, 1), case.num_heads) - scores = ( - torch.einsum("hd,hkd->hk", group_query.float(), keys.float()) - * module.scaling - ) - score_parts.append(scores) - value_parts.append(values.float()) - - scores = torch.cat(score_parts, dim=-1) - values = torch.cat(value_parts, dim=1) - probs = torch.softmax(scores, dim=-1) - out = torch.einsum("hk,hkd->hd", probs, values) - outputs.append(out.reshape(-1)) - - attn_output = torch.stack(outputs, dim=0).to(dtype) - return module.reconstruct_output(attn_output) - - -def _dual_chunk_sparse_fallback_indices_reference( - seq_len: int, max_count: int, device: torch.device -) -> torch.Tensor: - count = min(int(max_count), seq_len) - if count <= 0: - return torch.empty(0, dtype=torch.int64, device=device) - step = max(1, (seq_len + count - 1) // count) - return torch.arange(0, seq_len, step, dtype=torch.int64, device=device)[:count] - - -def _dual_chunk_sum_all_diagonal_matrix(mat: torch.Tensor) -> torch.Tensor: - h, n, m = mat.shape - zero_mat = torch.zeros((h, n, n), dtype=mat.dtype, device=mat.device) - mat_padded = torch.cat((zero_mat, mat, zero_mat), -1) - mat_strided = mat_padded.as_strided( - (1, n, n + m), (n * (2 * n + m), 2 * n + m + 1, 1) - ) - return torch.sum(mat_strided, 1)[:, 1:] - - -def _normalise_sparse_stage_indices( - indices: torch.Tensor, - counts: torch.Tensor, - *, - descending: bool, -) -> tuple[tuple[int, ...], ...]: - indices = indices.detach().reshape(1, counts.numel(), -1) - counts = counts.detach().to(torch.int64).cpu().tolist() - stage_indices = [] - for head_i, count in enumerate(counts): - head_indices = indices[0, head_i, :count].to(torch.int64) - head_indices = head_indices.sort(descending=descending).values - stage_indices.append(tuple(int(idx) for idx in head_indices.cpu().tolist())) - return tuple(stage_indices) - - -def _make_sparse_stage_selection( - stage: str, - q_len: int, - kv_len: int, - vertical_indices: list[torch.Tensor], - slash_indices: list[torch.Tensor], -) -> _DualChunkSparseStageSelection: - return _DualChunkSparseStageSelection( - stage=stage, - q_len=q_len, - kv_len=kv_len, - vertical_indices=tuple( - tuple(int(idx) for idx in head_indices.sort().values.cpu().tolist()) - for head_indices in vertical_indices - ), - slash_indices=tuple( - tuple( - int(idx) - for idx in head_indices.sort(descending=True).values.cpu().tolist() - ) - for head_indices in slash_indices - ), - ) - - -def _capture_sparse_stage_selection( - stage: str, - query: torch.Tensor, - key: torch.Tensor, - vertical_indices: torch.Tensor, - slash_indices: torch.Tensor, - vertical_counts: torch.Tensor | None, - slash_counts: torch.Tensor | None, -) -> _DualChunkSparseStageSelection: - assert vertical_counts is not None - assert slash_counts is not None - return _DualChunkSparseStageSelection( - stage=stage, - q_len=query.shape[2], - kv_len=key.shape[2], - vertical_indices=_normalise_sparse_stage_indices( - vertical_indices, vertical_counts, descending=False - ), - slash_indices=_normalise_sparse_stage_indices( - slash_indices, slash_counts, descending=True - ), - ) - - -def _dual_chunk_sparse_selection_reference( - fixture: DualChunkAttentionFixture, -) -> list[_DualChunkSparseStageSelection]: - """Reference DCA's content-aware top-k split and fallback indices. - - This covers the DCA-specific part of sparse prefill: vertical/slash top-k - selection, intra/succ/inter splitting, and empty-stage fallback. The lower - level 64x64 block conversion and sparse kernel math are covered separately. - """ - case = fixture.case - assert case.batch_size == 1 - assert case.prefix_lens == (0,) - assert case.num_heads == case.num_kv_heads - - module = fixture.reference_module - ( - q, - q_succ, - q_inter, - q_succ_critical, - q_inter_critical, - k, - _, - ) = module.project_dual_qkv(fixture.input_hidden) - q = q.view(-1, case.num_heads, module.head_dim) - q_succ = q_succ.view(-1, case.num_heads, module.head_dim) - q_inter = q_inter.view(-1, case.num_heads, module.head_dim) - q_succ_critical = q_succ_critical.view(-1, case.num_heads, module.head_dim) - q_inter_critical = q_inter_critical.view(-1, case.num_heads, module.head_dim) - k = k.view(-1, case.num_kv_heads, module.head_dim) - - config = DUAL_CHUNK_SPARSE_SUB_WINDOW_CONFIG - chunk_len = config["chunk_size"] - config["local_size"] - softmax_scale = module.scaling - scaling_factor = ( - 0.1 - * torch.log( - torch.tensor(case.seq_lens[0] / config["original_max_position_embeddings"]) - ) - + 1.0 - ).clamp(min=1) - softmax_scale *= float(scaling_factor.item()) - - head_config = config["sparse_attention_config"][0] - heads_vertical_size = [] - heads_slash_size = [] - for head_i in range(case.num_heads): - ty, vertical_size, slash_size, _ = head_config[str(head_i)] - assert ty == "vertical_and_slash" - if vertical_size == 30: - vertical_size += 100 - heads_vertical_size.append(vertical_size) - heads_slash_size.append(slash_size) - - selections = [] - k_length = k.shape[0] - begin = k_length - q.shape[0] - while begin < k_length: - prev_chunk_end_pos = (begin // chunk_len) * chunk_len - next_chunk_end_pos = prev_chunk_end_pos + chunk_len - end = min(next_chunk_end_pos, k_length) - qbegin = begin - (k_length - q.shape[0]) - qend = end - (k_length - q.shape[0]) - chunk_q_len = qend - qbegin - last_q_size = min(chunk_q_len, config["sparse_attention_last_q"]) - - q_states_intra = q[qbegin:qend] - k_states_intra = k[prev_chunk_end_pos:end] - qk_chunks = [ - (q_states_intra.transpose(0, 1)[:, -last_q_size:] * softmax_scale) - @ k_states_intra.permute(1, 2, 0) - ] - stage_kv_lens = {"intra": k_states_intra.size(0)} - - if prev_chunk_end_pos - chunk_len >= 0: - q_states_succ_critical = q_succ_critical[qbegin:qend] - k_states_succ = k[prev_chunk_end_pos - chunk_len : prev_chunk_end_pos] - qk_chunks.append( - ( - q_states_succ_critical.transpose(0, 1)[:, -last_q_size:] - * softmax_scale - ) - @ k_states_succ.permute(1, 2, 0) - ) - stage_kv_lens["succ"] = k_states_succ.size(0) - - if prev_chunk_end_pos - chunk_len * 2 >= 0: - q_states_inter_critical = q_inter_critical[qbegin:qend] - k_states_inter = k[: prev_chunk_end_pos - chunk_len] - qk_chunks.append( - ( - q_states_inter_critical.transpose(0, 1)[:, -last_q_size:] - * softmax_scale - ) - @ k_states_inter.permute(1, 2, 0) - ) - stage_kv_lens["inter"] = k_states_inter.size(0) - - qk = torch.cat(qk_chunks[::-1], dim=-1) - arange = torch.arange(last_q_size, device=q.device) - last_q_mask = arange[:, None] >= arange[None, :] - qk[:, :, -last_q_size:] = torch.where( - last_q_mask.unsqueeze(0), - qk[:, :, -last_q_size:], - -torch.inf, - ) - qk = torch.softmax(qk, dim=-1, dtype=torch.float32) - - vertical = qk.sum(-2, keepdim=True) - vertical[..., :30] = torch.inf - vertical = vertical.reshape(case.num_heads, -1) - max_vertical_topk = min(vertical.shape[-1], max(heads_vertical_size)) - max_slash_topk = max(heads_slash_size) - vertical_topk_buffer = torch.topk(vertical, max_vertical_topk, -1).indices - - slash_topk_buffer = torch.empty( - (case.num_heads, max_slash_topk), dtype=torch.int64, device=q.device - ) - current_vertical_size = [ - min(head_vertical_size, max_vertical_topk) - for head_vertical_size in heads_vertical_size - ] - current_slash_size = [] - for head_i in range(case.num_heads): - head_score = qk[head_i : head_i + 1, :, :] - slash_scores = _dual_chunk_sum_all_diagonal_matrix(head_score) - if head_score.size(1) != 1: - slash_scores = slash_scores[..., : -last_q_size + 1] - slash_scores[..., -100:] = torch.inf - - head_slash_size = min(heads_slash_size[head_i], vertical.size(-1)) - current_slash_size.append(head_slash_size) - slash_topk = torch.topk(slash_scores, head_slash_size, -1).indices - slash_topk_buffer[head_i, :head_slash_size] = slash_topk.reshape(-1) - - stage_vertical_indices = {stage: [] for stage in stage_kv_lens} - stage_slash_indices = {stage: [] for stage in stage_kv_lens} - for head_i in range(case.num_heads): - vertical_topk = vertical_topk_buffer[ - head_i, : current_vertical_size[head_i] - ] - slash_topk = slash_topk_buffer[head_i, : current_slash_size[head_i]] - - intra_vertical_indices = ( - vertical_topk[vertical_topk >= prev_chunk_end_pos] - prev_chunk_end_pos - ) - if intra_vertical_indices.nelement() == 0: - intra_vertical_indices = _dual_chunk_sparse_fallback_indices_reference( - stage_kv_lens["intra"], current_vertical_size[head_i], q.device - ) - intra_slash_indices = (qk.size(-1) - 1) - slash_topk[ - slash_topk >= prev_chunk_end_pos - ] - if intra_slash_indices.nelement() == 0: - intra_slash_indices = _dual_chunk_sparse_fallback_indices_reference( - stage_kv_lens["intra"], current_slash_size[head_i], q.device - ) - stage_vertical_indices["intra"].append(intra_vertical_indices) - stage_slash_indices["intra"].append(intra_slash_indices) - - if "succ" in stage_kv_lens: - succ_vertical_indices = vertical_topk[ - (vertical_topk < prev_chunk_end_pos) - & (vertical_topk >= prev_chunk_end_pos - chunk_len) - ] - (prev_chunk_end_pos - chunk_len) - if succ_vertical_indices.nelement() == 0: - succ_vertical_indices = ( - _dual_chunk_sparse_fallback_indices_reference( - stage_kv_lens["succ"], - current_vertical_size[head_i], - q.device, - ) - ) - succ_slash_indices = ( - prev_chunk_end_pos + chunk_q_len - 1 - ) - slash_topk[ - (slash_topk >= (prev_chunk_end_pos - chunk_len)) - & (slash_topk < (prev_chunk_end_pos + chunk_q_len)) - ] - if succ_slash_indices.nelement() == 0: - succ_slash_indices = _dual_chunk_sparse_fallback_indices_reference( - stage_kv_lens["succ"], - current_slash_size[head_i], - q.device, - ) - stage_vertical_indices["succ"].append(succ_vertical_indices) - stage_slash_indices["succ"].append(succ_slash_indices) - - if "inter" in stage_kv_lens: - inter_vertical_indices = vertical_topk[ - vertical_topk < prev_chunk_end_pos - chunk_len - ] - if inter_vertical_indices.nelement() == 0: - inter_vertical_indices = ( - _dual_chunk_sparse_fallback_indices_reference( - stage_kv_lens["inter"], - current_vertical_size[head_i], - q.device, - ) - ) - inter_slash_indices = ( - prev_chunk_end_pos - chunk_len + chunk_q_len - 1 - ) - slash_topk[ - slash_topk < (prev_chunk_end_pos - chunk_len + chunk_q_len) - ] - if inter_slash_indices.nelement() == 0: - inter_slash_indices = _dual_chunk_sparse_fallback_indices_reference( - stage_kv_lens["inter"], - current_slash_size[head_i], - q.device, - ) - stage_vertical_indices["inter"].append(inter_vertical_indices) - stage_slash_indices["inter"].append(inter_slash_indices) - - for stage in ("intra", "succ", "inter"): - if stage not in stage_kv_lens: - continue - selections.append( - _make_sparse_stage_selection( - stage, - chunk_q_len, - stage_kv_lens[stage], - stage_vertical_indices[stage], - stage_slash_indices[stage], - ) - ) - begin = end - - return selections - - -def _assert_sparse_stage_selections_match( - testcase, - actual: list[_DualChunkSparseStageSelection], - expected: list[_DualChunkSparseStageSelection], -) -> None: - testcase.assertEqual( - len(actual), - len(expected), - f"expected {len(expected)} sparse stage calls, got {len(actual)}", - ) - for index, (actual_stage, expected_stage) in enumerate(zip(actual, expected)): - testcase.assertEqual( - actual_stage, - expected_stage, - f"sparse stage selection mismatch at call {index}", - ) - - -def _run_dual_chunk_fixture_with_sparse_selection_capture( - fixture: DualChunkAttentionFixture, -) -> tuple[torch.Tensor, list[_DualChunkSparseStageSelection]]: - original_sparse_attention = _dual_chunk_backend._vertical_slash_sparse_attention - selections = [] - - def capture_sparse_attention( - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - vertical_indices: torch.Tensor, - slash_indices: torch.Tensor, - softmax_scale: float, - causal: bool = True, - stage: str = "intra", - block_size_M: int = 64, - block_size_N: int = 64, - vertical_indices_count: torch.Tensor | None = None, - slash_indices_count: torch.Tensor | None = None, - ): - selections.append( - _capture_sparse_stage_selection( - stage, - query, - key, - vertical_indices, - slash_indices, - vertical_indices_count, - slash_indices_count, - ) - ) - return original_sparse_attention( - query, - key, - value, - vertical_indices, - slash_indices, - softmax_scale, - causal=causal, - stage=stage, - block_size_M=block_size_M, - block_size_N=block_size_N, - vertical_indices_count=vertical_indices_count, - slash_indices_count=slash_indices_count, - ) - - try: - _dual_chunk_backend._vertical_slash_sparse_attention = capture_sparse_attention - output = run_dual_chunk_fixture_eager(fixture) - finally: - _dual_chunk_backend._vertical_slash_sparse_attention = original_sparse_attention - return output, selections - - -def _torch_sparse_attn_metadata_mask( - block_count: torch.Tensor, - block_offset: torch.Tensor, - column_count: torch.Tensor, - column_index: torch.Tensor, - q_len: int, - kv_len: int, - *, - block_size_m: int = 64, - block_size_n: int = 64, -) -> torch.Tensor: - batch_size, num_heads, num_rows = block_count.shape - mask = torch.zeros( - (batch_size, num_heads, q_len, kv_len), - dtype=torch.bool, - device=block_count.device, - ) - - for batch_i in range(batch_size): - for head_i in range(num_heads): - for row_i in range(num_rows): - row_start = row_i * block_size_m - row_end = min(row_start + block_size_m, q_len) - if row_start >= row_end: - continue - - for block_i in range(int(block_count[batch_i, head_i, row_i].item())): - col_start = int( - block_offset[batch_i, head_i, row_i, block_i].item() - ) - col_end = min(col_start + block_size_n, kv_len) - if 0 <= col_start < col_end: - mask[batch_i, head_i, row_start:row_end, col_start:col_end] = ( - True - ) - - col_count = int(column_count[batch_i, head_i, row_i].item()) - cols = column_index[batch_i, head_i, row_i, :col_count].to(torch.long) - cols = cols[(cols >= 0) & (cols < kv_len)] - if cols.numel() > 0: - mask[batch_i, head_i, row_start:row_end, cols] = True - - return mask - - -def _torch_sparse_attn_func( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - block_count: torch.Tensor, - block_offset: torch.Tensor, - column_count: torch.Tensor, - column_index: torch.Tensor, - dropout_p: float = 0.0, - softmax_scale: float | None = None, - causal: bool = False, - softcap: float = 0.0, - alibi_slopes: torch.Tensor | None = None, - deterministic: bool = False, - return_attn_probs: bool = False, - *, - return_softmax_lse: bool = False, - out: torch.Tensor | None = None, -): - assert dropout_p == 0.0 - assert softcap == 0.0 - assert alibi_slopes is None - assert not deterministic - assert not return_attn_probs - assert out is None - - if softmax_scale is None: - softmax_scale = q.shape[-1] ** -0.5 - - dtype = q.dtype - _, q_len, num_heads, _ = q.shape - kv_len = k.shape[1] - if k.shape[2] != num_heads: - group_size = num_heads // k.shape[2] - k = torch.repeat_interleave(k, group_size, dim=2) - v = torch.repeat_interleave(v, group_size, dim=2) - - sparse_mask = _torch_sparse_attn_metadata_mask( - block_count, block_offset, column_count, column_index, q_len, kv_len - ) - if causal: - q_pos = torch.arange(q_len, device=q.device) + (kv_len - q_len) - k_pos = torch.arange(kv_len, device=q.device) - sparse_mask &= k_pos.view(1, 1, 1, kv_len) <= q_pos.view(1, 1, q_len, 1) - - scores = torch.einsum("bqhd,bkhd->bhqk", q.float(), k.float()) * softmax_scale - scores = scores.masked_fill(~sparse_mask, -torch.inf) - softmax_lse = torch.logsumexp(scores, dim=-1) - valid_rows = sparse_mask.any(dim=-1) - probs = torch.softmax(scores, dim=-1) - probs = torch.where(valid_rows.unsqueeze(-1), probs, torch.zeros_like(probs)) - output = torch.einsum("bhqk,bkhd->bqhd", probs, v.float()).to(dtype) - - if return_softmax_lse: - return output, softmax_lse - return output - - -def _run_dual_chunk_fixture_with_torch_sparse_kernel( - fixture: DualChunkAttentionFixture, -) -> torch.Tensor: - original_sparse_attn_func = _dual_chunk_backend.sparse_attn_func - try: - _dual_chunk_backend.sparse_attn_func = _torch_sparse_attn_func - return run_dual_chunk_fixture_eager(fixture) - finally: - _dual_chunk_backend.sparse_attn_func = original_sparse_attn_func - - -def run_dual_chunk_attention_case( - testcase, - case: DualChunkAttentionCase, - *, - head_dim: int = DEFAULT_HEAD_DIM, - hidden_size: int = DEFAULT_HIDDEN_SIZE, - max_context_len: int = DEFAULT_MAX_CONTEXT_LEN, - dtype: torch.dtype = DEFAULT_DTYPE, - device: str = DEFAULT_DEVICE, - loc_layout: str = "shuffled_pages", -) -> None: - fixture = build_dual_chunk_attention_fixture( - testcase, - case, - head_dim=head_dim, - hidden_size=hidden_size, - max_context_len=max_context_len, - dtype=dtype, - device=device, - loc_layout=loc_layout, - ) - actual = run_dual_chunk_fixture_eager(fixture) - expected = expected_dual_chunk_fixture_output(fixture) - torch.testing.assert_close(actual, expected, atol=DENSE_ATOL, rtol=DENSE_RTOL) - - -def run_dual_chunk_sparse_attention_case( - testcase, - case: DualChunkAttentionCase, - *, - max_context_len: int = DEFAULT_MAX_CONTEXT_LEN, - dtype: torch.dtype = DEFAULT_DTYPE, - device: str = DEFAULT_DEVICE, -) -> None: - fixture = build_dual_chunk_attention_fixture( - testcase, - case, - # The local sparse FlashAttention build only includes head_dim=128. - head_dim=128, - hidden_size=128, - max_context_len=max_context_len, - dtype=dtype, - device=device, - dual_chunk_attention_config=DUAL_CHUNK_SPARSE_ALL_COLUMN_CONFIG, - ) - actual = run_dual_chunk_fixture_eager(fixture) - expected = expected_dual_chunk_fixture_output(fixture) - torch.testing.assert_close(actual, expected, atol=DENSE_ATOL, rtol=DENSE_RTOL) - - -def run_dual_chunk_sparse_threshold_gated_case( - testcase, - case: DualChunkAttentionCase, - *, - max_context_len: int = DEFAULT_MAX_CONTEXT_LEN, - dtype: torch.dtype = DEFAULT_DTYPE, - device: str = DEFAULT_DEVICE, -) -> None: - fixture = build_dual_chunk_attention_fixture( - testcase, - case, - head_dim=128, - hidden_size=128, - max_context_len=max_context_len, - dtype=dtype, - device=device, - dual_chunk_attention_config=DUAL_CHUNK_SPARSE_THRESHOLD_GATED_CONFIG, - ) - actual = run_dual_chunk_fixture_eager(fixture) - expected = expected_dual_chunk_fixture_output(fixture) - torch.testing.assert_close(actual, expected, atol=DENSE_ATOL, rtol=DENSE_RTOL) - - -def run_dual_chunk_sparse_sub_window_case( - testcase, - case: DualChunkAttentionCase, - *, - max_context_len: int = DEFAULT_MAX_CONTEXT_LEN, - dtype: torch.dtype = DEFAULT_DTYPE, - device: str = DEFAULT_DEVICE, -) -> None: - """Correctness test for genuine sub-context-window sparse pruning.""" - fixture = build_dual_chunk_attention_fixture( - testcase, - case, - head_dim=128, - hidden_size=128, - max_context_len=max_context_len, - dtype=dtype, - device=device, - dual_chunk_attention_config=DUAL_CHUNK_SPARSE_SUB_WINDOW_CONFIG, - ) - actual, actual_selections = _run_dual_chunk_fixture_with_sparse_selection_capture( - fixture - ) - expected_selections = _dual_chunk_sparse_selection_reference(fixture) - _assert_sparse_stage_selections_match( - testcase, actual_selections, expected_selections - ) - expected = _run_dual_chunk_fixture_with_torch_sparse_kernel(fixture) - torch.testing.assert_close(actual, expected, atol=DENSE_ATOL, rtol=DENSE_RTOL) - - -# --------------------------------------------------------------------------- -# Runner-mode helpers (mirror dense conventions; dual-chunk wraps RadixAttention -# so the K-write happens via `save_kv_cache=True` inside the backend forward). -# --------------------------------------------------------------------------- - - -def make_dual_chunk_case_with_prefix_lens( - case: DualChunkAttentionCase, - name: str, - prefix_lens: tuple[int, ...], -) -> DualChunkAttentionCase: - if case.forward_mode.is_decode(): - extend_lens: tuple[int, ...] = () - else: - base = case.extend_lens or (1,) - if len(prefix_lens) <= len(base): - extend_lens = base[: len(prefix_lens)] - else: - extend_lens = base + (base[-1],) * (len(prefix_lens) - len(base)) - return DualChunkAttentionCase( - name=name, - backend=case.backend, - forward_mode=case.forward_mode, - num_heads=case.num_heads, - num_kv_heads=case.num_kv_heads, - page_size=case.page_size, - prefix_lens=prefix_lens, - extend_lens=extend_lens, - ) - - -def dual_chunk_fixture_inputs( - fixture: DualChunkAttentionFixture, -) -> dict: - return { - "prefix_hidden": fixture.prefix_hidden, - "input_hidden": fixture.input_hidden, - } - - -def make_dual_chunk_random_inputs( - case: DualChunkAttentionCase, - fixture: DualChunkAttentionFixture, - *, - dtype: torch.dtype, - device: str, -) -> dict: - hidden_size = fixture.actual_module.hidden_size - prefix_hidden = [ - torch.randn(length, hidden_size, dtype=dtype, device=device) - for length in case.prefix_lens - ] - input_hidden = torch.randn( - case.num_input_tokens, hidden_size, dtype=dtype, device=device - ) - return {"prefix_hidden": prefix_hidden, "input_hidden": input_hidden} - - -def make_dual_chunk_replay_inputs( - case: DualChunkAttentionCase, - fixture: DualChunkAttentionFixture, - pad_prefix_lens: tuple[int, ...], - base_inputs: dict, - *, - dtype: torch.dtype, - device: str, -) -> dict: - """Pad the base inputs with random prefix/input hidden for the trailing - padding requests so the replay batch matches the capture-batch shape.""" - hidden_size = fixture.actual_module.hidden_size - pad_prefix_hidden = [ - torch.randn(length, hidden_size, dtype=dtype, device=device) - for length in pad_prefix_lens - ] - extra_input_tokens = case.num_input_tokens - base_inputs["input_hidden"].shape[0] - if extra_input_tokens < 0: - raise ValueError("padded case must have at least as many input tokens as base.") - pad_input_hidden = torch.randn( - extra_input_tokens, hidden_size, dtype=dtype, device=device - ) - return { - "prefix_hidden": base_inputs["prefix_hidden"] + pad_prefix_hidden, - "input_hidden": torch.cat( - [base_inputs["input_hidden"], pad_input_hidden], dim=0 - ), - } - - -def prepare_dual_chunk_runner_inputs( - fixture: DualChunkAttentionFixture, - case: DualChunkAttentionCase, - batch: ForwardBatch, - inputs: dict, - *, - max_context_len: int, -) -> None: - """Rebind inputs on the fixture, set `batch.orig_seq_lens` (dual-chunk - reads it during forward), and re-populate prefix K cache for the - (possibly re-shaped) case.""" - fixture.case = case - fixture.forward_batch = batch - fixture.prefix_hidden = inputs["prefix_hidden"] - fixture.input_hidden = inputs["input_hidden"] - _set_orig_seq_lens(batch, case) - _populate_prefix_kv( - fixture.actual_module, - case, - fixture.runner, - fixture.prefix_hidden, - max_context_len=max_context_len, - ) - - -def run_dual_chunk_forward( - fixture: DualChunkAttentionFixture, - batch: ForwardBatch, - inputs: dict, -) -> torch.Tensor: - return fixture.actual_module(inputs["input_hidden"], batch) - - -def expected_dual_chunk_output_from_inputs( - fixture: DualChunkAttentionFixture, - case: DualChunkAttentionCase, - inputs: dict, - state, -) -> torch.Tensor: - del state - return _dual_chunk_attention_reference( - fixture.reference_module, - case, - inputs["prefix_hidden"], - inputs["input_hidden"], - ) - - -def dual_chunk_attention_layers(fixture: DualChunkAttentionFixture) -> list: - return [fixture.actual_module.attn] - - -def _clone_dual_chunk_cache(fixture: DualChunkAttentionFixture): - """Snapshot the layer's K cache buffer. Dual-chunk writes K cache via - `set_kv_buffer` at decode time, so the capture forward's K write - persists into replay; the snapshot lets us roll it back.""" - layer_id = fixture.actual_module.attn.layer_id - kv_buf = fixture.runner.token_to_kv_pool.get_key_buffer(layer_id) - v_buf = fixture.runner.token_to_kv_pool.get_value_buffer(layer_id) - return (kv_buf.clone(), v_buf.clone()) - - -def _restore_dual_chunk_cache(fixture: DualChunkAttentionFixture, state) -> None: - layer_id = fixture.actual_module.attn.layer_id - fixture.runner.token_to_kv_pool.get_key_buffer(layer_id).copy_(state[0]) - fixture.runner.token_to_kv_pool.get_value_buffer(layer_id).copy_(state[1]) diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/cuda_graph_decode_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/cuda_graph_decode_runner.py index 605cef467..b702744d1 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/cuda_graph_decode_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/cuda_graph_decode_runner.py @@ -73,20 +73,6 @@ from ..attention_methods.dsv4_attention import ( run_dsv4_fixture_eager, run_dsv4_forward, ) -from ..attention_methods.dual_chunk_attention import ( - DualChunkAttentionCase, - _clone_dual_chunk_cache, - _restore_dual_chunk_cache, - build_dual_chunk_attention_fixture, - dual_chunk_fixture_inputs, - expected_dual_chunk_output_from_inputs, - make_dual_chunk_case_with_prefix_lens, - make_dual_chunk_random_inputs, - make_dual_chunk_replay_inputs, - prepare_dual_chunk_runner_inputs, - run_dual_chunk_fixture_eager, - run_dual_chunk_forward, -) from ..attention_methods.gdn_attention import DEFAULT_DEVICE as GDN_DEFAULT_DEVICE from ..attention_methods.gdn_attention import DEFAULT_DTYPE as GDN_DEFAULT_DTYPE from ..attention_methods.gdn_attention import ( @@ -891,58 +877,3 @@ def run_dsa_sparse_cuda_graph_decode_case( dtype=dtype, device=device, ) - - -def run_dual_chunk_cuda_graph_decode_case( - testcase, - case: DualChunkAttentionCase, - *, - head_dim: int = DEFAULT_HEAD_DIM, - hidden_size: int = DEFAULT_HIDDEN_SIZE, - max_context_len: int = DENSE_DEFAULT_MAX_CONTEXT_LEN, - dtype: torch.dtype = DENSE_DEFAULT_DTYPE, - device: str = DENSE_DEFAULT_DEVICE, - cuda_graph_capture_batch_size: int | None = None, -): - """Dual-chunk CUDA-graph decode replay. Decode reads cached K/V (set - by `set_kv_buffer` inside `forward_decode`) so the capture/replay - contract is the same shape as dense attention. The - `_clone_dual_chunk_cache` / `_restore_dual_chunk_cache` hooks snapshot - both K and V buffers so the capture forward's writes don't bleed into - replay state.""" - if not case.forward_mode.is_decode(): - raise ValueError("run_dual_chunk_cuda_graph_decode_case expects a DECODE case.") - capture_batch_size = cuda_graph_capture_batch_size or case.batch_size - adapter = CudaGraphDecodeAdapter( - build_fixture=build_dual_chunk_attention_fixture, - make_case=make_dual_chunk_case_with_prefix_lens, - make_forward_batch=_make_dense_forward_batch, - fixture_inputs=dual_chunk_fixture_inputs, - make_capture_inputs=make_dual_chunk_random_inputs, - make_replay_inputs=make_dual_chunk_replay_inputs, - prepare_inputs=prepare_dual_chunk_runner_inputs, - run_eager=run_dual_chunk_fixture_eager, - run_forward=run_dual_chunk_forward, - expected_output=expected_dual_chunk_output_from_inputs, - clone_state=_clone_dual_chunk_cache, - restore_state=_restore_dual_chunk_cache, - allow_padding=True, - atol=DENSE_ATOL, - rtol=DENSE_RTOL, - ) - _run_cuda_graph_decode_case( - testcase, - case, - adapter=adapter, - build_kwargs=dict( - head_dim=head_dim, - hidden_size=hidden_size, - max_context_len=max_context_len, - dtype=dtype, - device=device, - ), - capture_batch_size=capture_batch_size, - max_context_len=max_context_len, - dtype=dtype, - device=device, - ) diff --git a/test/registered/attention/unittests/KNOWN_FAILURES.md b/test/registered/attention/unittests/KNOWN_FAILURES.md index 04d599f67..67a7b7742 100644 --- a/test/registered/attention/unittests/KNOWN_FAILURES.md +++ b/test/registered/attention/unittests/KNOWN_FAILURES.md @@ -34,29 +34,6 @@ hardware/version X). # A. Container re-image required -## A.1. `flash_attn` SM10.x wheel missing - -**Affected**: `dual_chunk/test_dual_chunk_flash_attn.py` (entire class — 5 -test methods, ~18 subtests) - -**Symptom on GB300**: -``` -ImportError: cannot import name 'flash_attn_varlen_func' from 'flash_attn' -``` - -**Root cause**: `DualChunkFlashAttentionBackend` calls `flash_attn_varlen_func` -via `sglang.kernels.ops.attention.flash_attention`. On SM 8.x / 9.x that resolves to -sgl-kernel's FA3 build (works on H200). On other SMs, the JIT kernel falls -back to the upstream `flash_attn` (FA2) wheel — but the -`lmsysorg/sglang:nightly-dev-cu13` container's `flash_attn` package on -SM10.x is missing `flash_attn_varlen_func`. - -**Gate**: `_dual_chunk_fa_supported()` in -`dual_chunk/test_dual_chunk_flash_attn.py` skips the whole class on the -fallback-broken path. Hopper passes through unchanged. - -**Fix**: Re-image with an SM10.x-compiled `flash_attn` wheel. - ## A.2. tilelang `wait_wgmma` template missing on SM10.x **Affected**: @@ -89,7 +66,6 @@ so that "skipped: ..." results have a quick lookup. | Backend | Required SM | Gate location | Error if unguarded | |---|---|---|---| -| `cutlass_mla` | exactly SM 10.0 (B200) | `mla/test_cutlass_mla.py::_supported` | `cutlass_mla_decode is only supported on compute capability 10.0, but found sm version 103` | | `flashmla` decode/verify | SM 9.0 (Hopper) only | `mla/test_flashmla.py:_DECODE_REQUIRES_SM90A` | `Dense decode MLA is only supported on SM90a architecture` | | `trtllm_mla` | SM 12.0a / 12.1a | `mla/test_trtllm_mla.py::_supported` | FlashInfer XQA MLA dispatch reject | | `tokenspeed_mla` | SM ≥ 10.0 + FP8 KV + pkg | `mla/test_tokenspeed_mla.py::_supported` | `tokenspeed_mla` import or kernel dispatch | @@ -99,7 +75,7 @@ so that "skipped: ..." results have a quick lookup. | `fa3` (non-MLA) | SM 80 or SM 90 | `_is_fa3_supported` in `flash_attention_v3.py` | `attention_registry.py:177-180` reject | **SM10.3 vs SM10.0**: GB300 is SM10.3. Gates that require exactly SM10.0 -(cutlass_mla, dsa trtllm) intentionally skip on GB300 because the kernel +(dsa trtllm) intentionally skip on GB300 because the kernel binaries in the container aren't compiled for sm_103. Flip the gates to `major == 10` (drop the `minor == 0`) once GB300-compiled binaries land. @@ -133,13 +109,7 @@ methods that record each backend's failure mode inline as | `mla/test_flashmla.py::test_layout_robustness_cases` (extend) | `non_monotonic_extend` | FlashMLA extend raises CUDA illegal memory access. | | `mla/test_flashmla.py::test_layout_robustness_cases` (decode) | `interleaved_pages` | FlashMLA decode raises `shape '[-1, 64, 1, 32]' is invalid for input of size N`. | -### Dual-chunk - -| Test | Layout | Root cause | -|---|---|---| -| `dual_chunk/test_dual_chunk_flash_attn.py::test_layout_robustness_cases` (extend) | `non_monotonic_extend` | `_dual_chunk_flash_attn_prefill_func` uses `cu_seqlens_*` indexing into contiguous K slots (`dual_chunk_flashattention_backend.py:834+`); scattered extend-token slots break that contiguity. | - -**Total**: 9 layout-handling production bugs documented. +**Total**: 8 layout-handling production bugs documented. ## C.2. Speculative-mode rejects @@ -205,7 +175,6 @@ fail at backend init. | Backend | Required page size(s) | Citation | |---|---|---| | FlashMLA | `64` only | `server_args.py:2767-2770` | -| Cutlass MLA | `128` only | `server_args.py:2776-2779`, `cutlass_mla_backend.py:31` | | TRT-LLM MLA | `{32, 64}` | `server_args.py:2790-2794` | | Tokenspeed MLA | `{32, 64}` | `server_args.py:2809-2813`, `tokenspeed_mla_backend.py:111-113` | | TRT-LLM MHA | `{16, 32, 64}` | `server_args.py:2849-2853` | @@ -241,12 +210,9 @@ fail at backend init. | Test file | Failure type | Section | |---|---|---| -| `dual_chunk/test_dual_chunk_flash_attn.py` | Container: `flash_attn` SM10.x wheel | §A.1 | -| `dual_chunk/test_dual_chunk_flash_attn.py::test_layout_robustness_cases` (non_monotonic_extend) | Layout-handling bug | §C.1 | | `dsa/test_dsa.py::test_sparse_tilelang_*` | Container: tilelang `wait_wgmma` | §A.2 | | `dsa/test_dsa.py::test_sparse_*_impl_variants` (tilelang row) | Container: tilelang `wait_wgmma` | §A.2 | | `dsa/test_dsa.py::test_sparse_*_impl_variants` (fa3 / trtllm rows) | Hardware gate | §B | -| `mla/test_cutlass_mla.py` (all) | Hardware gate (SM 10.0 exactly) | §B | | `mla/test_flashmla.py` (DECODE/verify subtests) | Hardware gate (SM 9.0 Hopper) | §B | | `mla/test_flashinfer.py::test_runner_mode_eagle_draft_cuda_graph_runner_cases` | Backend bug gated on SM≥10 | §C.3 | | `mla/test_flashinfer.py::test_layout_robustness_cases` | Layout-handling bug | §C.1 | diff --git a/test/registered/attention/unittests/dual_chunk/README.md b/test/registered/attention/unittests/dual_chunk/README.md deleted file mode 100644 index 18e6231ac..000000000 --- a/test/registered/attention/unittests/dual_chunk/README.md +++ /dev/null @@ -1,90 +0,0 @@ -# Dual-Chunk Attention Capability Matrix - -This folder covers dual-chunk attention tests. `dual_chunk_flash_attn` is not -a dense backend swap: it expects a packed five-way query projection (`query`, -`succ`, `inter`, and critical variants), so the dense Q/K/V harness is -structurally wrong for this method. The single attention backend here is -`dual_chunk_flash_attn`; the rows below distinguish kernel-path modes. - -## Coverage Matrix - -Columns are runner modes; rows are kernel-path modes of the single -`dual_chunk_flash_attn` backend. Cells use: -- **✓ \** — exercised, with the config variants listed in the cell -- **—** — not applicable / not exercised -- **blocked: \** — production-unsupported, not a follow-up -- **deferred: \** — could land later, currently disabled - -| Kernel path | Eager Phase 2 | CG decode | PCG extend | BCG extend | Verify eager | Verify CG | DE eager | DE CG | DE-V2 CG | EAGLE-draft runner | EAGLE-DE runner | FKVMTP runner | -|---|---|---|---|---|---|---|---|---|---|---|---|---| -| Non-sparse | ✓ first-window, successor-chunk, inter-chunk extend/decode layouts + GQA decode | deferred: graph metadata for dual-chunk not scoped | deferred | deferred | blocked: `init_forward_metadata` asserts `is_prefill() or is_decode()` (`dual_chunk_flashattention_backend.py:179`); `TARGET_VERIFY` falls under `is_prefill()` but the wrapper hasn't been wired through | deferred | deferred | deferred | blocked: `DRAFT_EXTEND_V2` excluded from `is_prefill()` alias (see Production-Unsupported below) | deferred | deferred | — | -| Sparse all-column (`vertical_size`/`slash_size` chosen so every key in the first chunk is selected) | ✓ single-request first-chunk, multi-request first-chunk, page-boundary first-chunk | — | — | — | blocked: same `is_prefill` assertion | — | — | — | blocked: same | — | — | — | -| Sparse sub-window (`vertical_size=4`, `slash_size=4`, `seq_len=128`) | ✓ independent DCA top-k/split/fallback reference + torch sparse-kernel reference | — | — | — | — | — | — | — | — | — | — | — | -| Threshold-gated sparse (`sparse_attention_threshold=100`, seq_len=16 → gate disables sparse, falls back to dense) | ✓ verifies `current_orig_seq_len > threshold` gate semantics | — | — | — | — | — | — | — | — | — | — | — | - -## Input And Config Coverage - -- Page size 1 extend, exact-page extend, page-boundary crossing extend, and - ragged extend batches. -- Decode page-boundary coverage and GQA decode coverage. -- Successor-chunk and inter-chunk extend/decode layouts where `query_succ` - and `query_inter` are active and use independent projection weights. -- Sparse all-column prefill uses `head_dim=128` to match the local sparse - FlashAttention build and selects every column in the first chunk - (≤16 tokens) so the dense reference remains valid. -- Multi-request sparse and page-boundary sparse variants exercise per-request - `cu_seqlens_*` slicing inside `_dual_chunk_flash_attn_prefill_func`. -- Sub-window sparse prefill uses `vertical_size=4`, `slash_size=4`, and - `seq_len=128` to verify the DCA-specific content-aware top-k split and - empty-stage fallback against an independent reference, then verifies the - sparse output against a torch sparse-kernel reference that consumes the - production block/column metadata. -- Threshold-gated sparse uses `sparse_attention_threshold=100` so a 16-token - prompt bypasses the sparse kernel and falls through to the dense chunk - flash path, exercising the gate semantics in the wrapper. - -## Container Gate (SM10.x) - -`DualChunkFlashAttentionBackend` calls `flash_attn_varlen_func` via -`sglang.kernels.ops.attention.flash_attention`. On SM8.x / SM9.x that resolves to sgl-kernel's -FA3 build; on SM != {8, 9} (notably SM10.x / GB300) the JIT kernel falls back -to the upstream `flash_attn` (FA2) wheel, which the -`lmsysorg/sglang:nightly-dev-cu13` container ships without an SM10.x-compiled -`flash_attn_varlen_func`. `test_dual_chunk_flash_attn.py` probes both paths at -module import: if FA3 is supported (`major in {8, 9}`) it runs unconditionally; -otherwise it tries `from flash_attn import flash_attn_varlen_func` and skips the -whole class with the documented reason if the symbol is missing. Re-image with -an SM10.x-compiled flash_attn wheel to clear; no test-code change needed. - -See `KNOWN_FAILURES.md` §1 for the full root cause + fix. - -## Production-Unsupported - -- **Non-prefill / non-decode forward modes** — - `dual_chunk_flashattention_backend.py:179` asserts - `forward_mode.is_prefill() or forward_mode.is_decode()`. `is_prefill()` - aliases to `is_extend()` (`forward_batch_info.py:103-104`) and covers - `EXTEND` / `MIXED` / `DRAFT_EXTEND` / `TARGET_VERIFY` / `SPLIT_PREFILL` / - `DLLM_EXTEND`, but `DRAFT_EXTEND_V2` is excluded by default. So - `DRAFT_EXTEND_V2` is structurally unreachable for `dual_chunk_flash_attn`. -- **Non-causal / windowed-attention requests** — `forward_extend` raises - `ValueError("Dual Chunk Attention does not support causal=False")` - (`dual_chunk_flashattention_backend.py:698`) and - `ValueError("Dual Chunk Attention does not support window_size")` - (`dual_chunk_flashattention_backend.py:700`). -- **Sparse mode `chunk_len % block_size != 0`** — raises - `ValueError("chunk_len must be divisible by block_size.")` - (`dual_chunk_flashattention_backend.py:860, 1491`). The current fixture - picks divisible values. -- **Unsupported `head_dim`** — only `head_dim in {16, 32, 64, 128, 256, 512}` - is accepted (`dual_chunk_flashattention_backend.py:1611`). - -## Next Work - -- Populate CUDA graph and PCG/BCG runner metadata after eager non-sparse - coverage is stable across more chunk layouts. -- **Broaden sub-window sparse coverage** — the current regression case covers - `prefix_lens=(0,)`, `extend_lens=(128,)`, and no GQA. Add - multi-request batches, nonzero prefixes, GQA, and more sparse config variants - once those paths need explicit sparse pruning coverage. The 64x64 - vertical/slash converter remains covered at the sgl-kernel layer. diff --git a/test/registered/attention/unittests/dual_chunk/__init__.py b/test/registered/attention/unittests/dual_chunk/__init__.py deleted file mode 100644 index 67141201d..000000000 --- a/test/registered/attention/unittests/dual_chunk/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Dual-chunk attention unit-test package.""" diff --git a/test/registered/attention/unittests/dual_chunk/test_dual_chunk_flash_attn.py b/test/registered/attention/unittests/dual_chunk/test_dual_chunk_flash_attn.py deleted file mode 100644 index 230aa4dda..000000000 --- a/test/registered/attention/unittests/dual_chunk/test_dual_chunk_flash_attn.py +++ /dev/null @@ -1,177 +0,0 @@ -import unittest - -import torch - -from sglang.srt.model_executor.forward_batch_info import ForwardMode -from sglang.test.kits.attention_unittest.attention_methods.dual_chunk_attention import ( - DualChunkAttentionCase, - make_dual_chunk_cases, - make_dual_chunk_sparse_cases, - make_dual_chunk_sparse_sub_window_cases, - make_dual_chunk_sparse_threshold_gated_cases, - run_dual_chunk_attention_case, - run_dual_chunk_sparse_attention_case, - run_dual_chunk_sparse_sub_window_case, - run_dual_chunk_sparse_threshold_gated_case, -) -from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner import ( - run_dual_chunk_cuda_graph_decode_case, -) -from sglang.test.test_utils import CustomTestCase - - -# Container gate (KNOWN_FAILURES.md §1): `DualChunkFlashAttentionBackend` calls -# `flash_attn_varlen_func` on every forward via -# `sglang.kernels.ops.attention.flash_attention`. On SM8x/SM9x, that resolves to sgl-kernel's -# FA3 build (which works). On SM != {8, 9} (notably SM10.3 / GB300), the JIT -# kernel falls back to the upstream `flash_attn` (FA2) wheel — but the -# `lmsysorg/sglang:nightly-dev-cu13` container's `flash_attn` package ships -# without `flash_attn_varlen_func` on SM10.x, so every dual-chunk forward -# fails at import time inside the fallback. Skip the whole suite only when -# that fallback path is actually broken (not on Hopper, where we never enter it). -# Re-image the container with an SM10.3-compiled flash_attn wheel to clear. -def _dual_chunk_fa_supported() -> tuple[bool, str]: - if not torch.cuda.is_available(): - return False, "CUDA is required" - major, _minor = torch.cuda.get_device_capability() - # FA3 path is taken when sm major is 8 or 9 (see - # `sglang.kernels.ops.attention.flash_attention_v3._is_fa3_supported`). On that path - # the upstream `flash_attn` fallback is never invoked. - if major in (8, 9): - return True, "" - # Otherwise (sm 7.x or sm >= 10.x) the JIT kernel falls back to upstream - # `flash_attn.flash_attn_varlen_func`. Probe it; if missing, skip. - try: - from flash_attn import ( # noqa: F401 - flash_attn_varlen_func as _flash_attn_varlen_func, - ) - - return True, "" - except ImportError as exc: - return False, ( - f"flash_attn_varlen_func is not available in upstream `flash_attn` " - f"(SM{major}.x JIT-kernel fallback): {exc}. " - f"Re-image the container with an SM{major}.x-compiled flash_attn wheel." - ) - - -_DUAL_CHUNK_FLASH_ATTN_AVAILABLE, _DUAL_CHUNK_SKIP_REASON = _dual_chunk_fa_supported() - - -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci(est_time=10, stage="base-b", runner_config="4-gpu-b200") -register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-large") - - -@unittest.skipIf(not torch.cuda.is_available(), "CUDA is required") -@unittest.skipIf(not _DUAL_CHUNK_FLASH_ATTN_AVAILABLE, _DUAL_CHUNK_SKIP_REASON) -class TestDualChunkFlashAttentionBackendCorrectness(CustomTestCase): - CASES = make_dual_chunk_cases("dual_chunk_flash_attn") - SPARSE_CASES = make_dual_chunk_sparse_cases("dual_chunk_flash_attn") - SPARSE_THRESHOLD_GATED_CASES = make_dual_chunk_sparse_threshold_gated_cases( - "dual_chunk_flash_attn" - ) - SPARSE_SUB_WINDOW_CASES = make_dual_chunk_sparse_sub_window_cases( - "dual_chunk_flash_attn" - ) - # Replay prefix_lens must each be >= capture_prefix_len (= fill-value - 1). - # Dual-chunk's `get_cuda_graph_seq_len_fill_value()` returns 1, so capture - # uses prefix=0. We pick a 3-request batch with varied lengths to exercise - # both the page-boundary and within-page slots. - CUDA_GRAPH_DECODE_CASES = ( - DualChunkAttentionCase( - name="runner_cuda_graph_dual_chunk_decode_page_boundary", - backend="dual_chunk_flash_attn", - forward_mode=ForwardMode.DECODE, - num_heads=4, - num_kv_heads=4, - page_size=16, - prefix_lens=(14, 15, 16), - ), - ) - - def test_projected_dual_chunk_attention_cases(self): - for case in self.CASES: - with self.subTest(case=case.name, backend=case.backend): - run_dual_chunk_attention_case(self, case) - - def test_sparse_dual_chunk_attention_cases(self): - for case in self.SPARSE_CASES: - with self.subTest(case=case.name, backend=case.backend): - run_dual_chunk_sparse_attention_case(self, case) - - def test_sparse_dual_chunk_threshold_gated_cases(self): - for case in self.SPARSE_THRESHOLD_GATED_CASES: - with self.subTest(case=case.name, backend=case.backend): - run_dual_chunk_sparse_threshold_gated_case(self, case) - - def test_sparse_dual_chunk_sub_window_cases(self): - for case in self.SPARSE_SUB_WINDOW_CASES: - with self.subTest(case=case.name, backend=case.backend): - run_dual_chunk_sparse_sub_window_case(self, case) - - def test_runner_mode_cuda_graph_decode_cases(self): - for case in self.CUDA_GRAPH_DECODE_CASES: - with self.subTest(case=case.name, backend=case.backend): - run_dual_chunk_cuda_graph_decode_case(self, case) - - # Layout-robustness. See dense/test_triton.py for the rationale. - # dual_chunk_flash_attn EXTEND fails on non_monotonic_extend with - # ~67% mismatch and max abs diff ~1.1. The dual-chunk prefill path - # uses `cu_seqlens_*` indexing into a contiguous K layout - # (see `_dual_chunk_flash_attn_prefill_func` in - # dual_chunk_flashattention_backend.py:834+), which assumes K for - # the new extend tokens is laid out contiguously in - # `[begin, end)` slot order. Scattering extend-token slots within a - # request breaks that contiguity. Documented as a known production - # limitation. - LAYOUT_ROBUSTNESS_CASES = ( - DualChunkAttentionCase( - name="layout_dual_chunk_extend_two_request", - backend="dual_chunk_flash_attn", - forward_mode=ForwardMode.EXTEND, - num_heads=4, - num_kv_heads=4, - page_size=16, - prefix_lens=(0, 0), - extend_lens=(16, 32), - ), - DualChunkAttentionCase( - name="layout_dual_chunk_decode_page_boundary", - backend="dual_chunk_flash_attn", - forward_mode=ForwardMode.DECODE, - num_heads=4, - num_kv_heads=4, - page_size=16, - prefix_lens=(14, 15, 16), - ), - ) - LAYOUT_KNOWN_FAILURES = { - ("layout_dual_chunk_extend_two_request", "non_monotonic_extend"): ( - "dual_chunk_flash_attn prefill uses cu_seqlens_* indexing " - "into contiguous K slots within an extend " - "(`_dual_chunk_flash_attn_prefill_func` in " - "dual_chunk_flashattention_backend.py:834+); scattered " - "extend-token slots break that contiguity." - ), - } - - def test_layout_robustness_cases(self): - for case in self.LAYOUT_ROBUSTNESS_CASES: - for layout in ("interleaved_pages", "non_monotonic_extend"): - if layout == "non_monotonic_extend" and case.forward_mode.is_decode(): - continue - reason = self.LAYOUT_KNOWN_FAILURES.get((case.name, layout)) - if reason is not None: - print( - f"[layout-known-failure] {case.name} x {layout}: {reason}", - flush=True, - ) - continue - with self.subTest(case=case.name, layout=layout): - run_dual_chunk_attention_case(self, case, loc_layout=layout) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/attention/unittests/mla/README.md b/test/registered/attention/unittests/mla/README.md index eedb92f4e..e4a08cc22 100644 --- a/test/registered/attention/unittests/mla/README.md +++ b/test/registered/attention/unittests/mla/README.md @@ -20,7 +20,6 @@ Columns are runner modes; rows are attention backends. Cells use: | `triton` | ✓ 10 input layouts (page 1/16/32, prefix/decode edges) | ✓ MLA decode page-boundary | ✓ ragged page-boundary extend | ✓ ragged page-boundary extend | ✓ EAGLE chain (topk=1) | ✓ EAGLE tree (topk=2) | — (V1 DE not enabled for Triton MLA; Triton uses V2 path) | — | ✓ fixed-tokens-per-req | ✓ chain (topk=1) + tree (topk=2) | ✓ via `DRAFT_EXTEND_V2` graph runner | — (no FKVMTP wiring for MLA) | | `flashinfer` | ✓ 10 input layouts with DeepSeek-like `kv_lora_rank=512`, `qk_rope_head_dim=64` | ✓ MLA decode page-boundary | ✓ ragged page-boundary extend | ✓ ragged page-boundary extend | ✓ EAGLE chain (topk=1) | ✓ EAGLE chain (topk=1) | ✓ EAGLE ragged-accept | ✓ EAGLE ragged-accept | blocked: `is_draft_extend()` default `include_v2=False` (`flashinfer_mla_backend.py:432,501,454-455,512`) | ✓ chain (topk=1) only — tree blocked by `topk=1` reject (`flashinfer_mla_backend.py:910-913`) | ✓ EAGLE ragged-accept (V1) | — (no FKVMTP wiring for MLA) | | `flashmla` | ✓ FlashMLA-compatible page-size-64 cases (zero-prefix exact page, input page edges 63/64/65, prefix exact page, total exact page, cross page, ragged, decode page-boundary, decode bsz=1 nonzero prefix) | ✓ page-size-64 decode page-boundary | ✓ ragged page-boundary extend | ✓ ragged page-boundary extend | ✓ EAGLE chain (topk=1) | ✓ EAGLE chain (topk=1) | ✓ EAGLE ragged-accept | deferred: parent FlashInfer-MLA capture path expects 1D `cuda_graph_kv_indices`, FlashMLA allocates 2D `[max_bs, (max_context+PAGE_SIZE)//PAGE_SIZE]` (`flashmla_backend.py:347-348` + parent `init_forward_metadata_capture_cuda_graph`) | — (FlashMLA does not implement V2) | ✓ chain (topk=1) only — tree blocked by `topk=1` reject (`flashmla_backend.py:555-558`) | — (DE CG deferred above) | — | -| `cutlass_mla` | skip:hw — needs SM 10.0+ (Blackwell); current 1 case uses `ForwardMode.EXTEND` but `CutlassMLABackend` only overrides `forward_decode` (`cutlass_mla_backend.py:226`) and falls through to FlashInfer MLA for other modes → **case should be DECODE**; PAGE_SIZE fixed at 128 (`cutlass_mla_backend.py:31`) | — (decode-only backend; no extend/CG) | — | — | blocked: tree via `topk=1` reject inherited from FlashInfer MLA parent | — | — | — | — | — | — | — | | `trtllm_mla` | skip:hw — needs SM 12.0a / 12.1a (`is_sm120_supported`) | — | — | — | blocked: `topk=1` only (`trtllm_mla_backend.py:1223-1229` inherits from FlashInfer MLA) | — | — | — | — | — | — | — | | `tokenspeed_mla` | skip:hw — needs `find_spec("tokenspeed_mla")`, SM 10.0+, and `kv_cache_dtype=fp8_e4m3` (`server_args.py:2814-2818`); current MLA fixture does not emit FP8 KV cache | — | — | — | blocked: `topk=1` only (`tokenspeed_mla_backend.py:341-347` inherits from TRT-LLM MLA) | — | — | — | — | — | — | — | @@ -60,17 +59,12 @@ multi-step draft backends and cannot ever appear at runtime. - **Tokenspeed MLA tree verify / draft-extend with `topk > 1`** — `TokenspeedMLAMultiStepDraftBackend` inherits from `TRTLLMMLAMultiStepDraftBackend` (`tokenspeed_mla_backend.py:341-347`). -- **Cutlass MLA extend / verify / draft-extend** — `CutlassMLABackend` only - overrides `forward_decode` (`cutlass_mla_backend.py:226`) and only handles - `is_decode_or_idle` in `init_forward_metadata*` (`cutlass_mla_backend.py:86, - 156, 197`). Anything else falls through to FlashInfer MLA. - **FlashInfer-MLA `DRAFT_EXTEND_V2` graph capture/replay** — `flashinfer_mla_backend.py:432,501` only route through `is_draft_extend()` (default `include_v2=False`); `else: raise ValueError("Invalid mode")` at `flashinfer_mla_backend.py:454-455,512`. - **All MLA backends fixed page size** — FlashMLA forces `page_size=64`, - Cutlass MLA forces `page_size=128`, TRT-LLM MLA and Tokenspeed MLA force - `page_size in {32, 64}`. + TRT-LLM MLA and Tokenspeed MLA force `page_size in {32, 64}`. ## Backend Container Gate (SM10.x) @@ -91,9 +85,6 @@ See `KNOWN_FAILURES.md` §3 for the full root cause + fix. override capture/replay in `FlashMLABackend` to use its 2D layout, or allocate both parent-style 1D and FlashMLA-style 2D buffers and route `DRAFT_EXTEND` to the parent path). -- Switch `mla/test_cutlass_mla.py` to `ForwardMode.DECODE` so it actually - exercises `CutlassMLABackend.forward_decode` instead of falling through to - FlashInfer MLA when SM 10.0+ is available. -- Add hardware-gated tests for `cutlass_mla`, `trtllm_mla`, and `tokenspeed_mla` +- Add hardware-gated tests for `trtllm_mla` and `tokenspeed_mla` decode (chain spec only) when the appropriate hardware/KV dtype fixtures are available. diff --git a/test/registered/attention/unittests/mla/test_cutlass_mla.py b/test/registered/attention/unittests/mla/test_cutlass_mla.py deleted file mode 100644 index 7587be457..000000000 --- a/test/registered/attention/unittests/mla/test_cutlass_mla.py +++ /dev/null @@ -1,98 +0,0 @@ -import unittest - -import torch - -from sglang.srt.model_executor.forward_batch_info import ForwardMode -from sglang.test.kits.attention_unittest.attention_methods.mla_attention import ( - MLAAttentionCase, - run_mla_attention_case, -) -from sglang.test.test_utils import CustomTestCase - -# Cutlass MLA requires exactly Blackwell SM 10.0. The sgl-kernel -# `cutlass_mla_decode` checks `sm_version == 100` (major*10+minor), so -# SM 10.3 (GB300) reports sm_version=103 and is rejected by the kernel. -# PAGE_SIZE is fixed to 128 in the backend. -_REQUIRED_SM_MAJOR = 10 -_REQUIRED_SM_MINOR = 0 - -MLA_SHAPE_KWARGS = dict( - kv_lora_rank=512, - qk_rope_head_dim=64, - hidden_size=1024, - max_context_len=256, -) - - -def _supported() -> tuple[bool, str]: - if not torch.cuda.is_available(): - return False, "CUDA is required" - major, minor = torch.cuda.get_device_capability() - if major != _REQUIRED_SM_MAJOR or minor != _REQUIRED_SM_MINOR: - return ( - False, - f"cutlass_mla requires exactly SM {_REQUIRED_SM_MAJOR}.{_REQUIRED_SM_MINOR} " - f"(B200 Blackwell); got SM {major}.{minor}", - ) - return True, "" - - -_SUPPORTED, _SKIP_REASON = _supported() - - -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci(est_time=12, stage="base-b", runner_config="4-gpu-b200") -register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-large") - - -@unittest.skipIf(not _SUPPORTED, _SKIP_REASON) -class TestCutlassMLAAttentionBackendCorrectness(CustomTestCase): - # CutlassMLABackend only overrides `forward_decode`; EXTEND falls through - # to the FlashInferMLAAttnBackend parent and bypasses cutlass code - # entirely. Use DECODE so the test actually exercises the cutlass kernel - # on Blackwell. Page size is fixed to PAGE_SIZE=128 (server_args.py - # forces this for cutlass_mla). - CASES = ( - MLAAttentionCase( - name="mla_decode_cutlass_page_boundary", - backend="cutlass_mla", - forward_mode=ForwardMode.DECODE, - num_heads=4, - page_size=128, - prefix_lens=(126, 127, 128), - ), - MLAAttentionCase( - name="mla_decode_cutlass_bsz1_nonzero_prefix", - backend="cutlass_mla", - forward_mode=ForwardMode.DECODE, - num_heads=4, - page_size=128, - prefix_lens=(63,), - ), - MLAAttentionCase( - name="mla_decode_cutlass_above_page", - backend="cutlass_mla", - forward_mode=ForwardMode.DECODE, - num_heads=4, - page_size=128, - prefix_lens=(128, 129, 130), - ), - MLAAttentionCase( - name="mla_decode_cutlass_multi_page", - backend="cutlass_mla", - forward_mode=ForwardMode.DECODE, - num_heads=4, - page_size=128, - prefix_lens=(127, 200, 255), - ), - ) - - def test_projected_mla_attention_cases(self): - for case in self.CASES: - with self.subTest(case=case.name, backend=case.backend): - run_mla_attention_case(self, case, **MLA_SHAPE_KWARGS) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/kernels/benchmark/quantization/bench_awq_dequantize.py b/test/registered/kernels/benchmark/quantization/bench_awq_dequantize.py index 6e547a4fb..35ca91b65 100644 --- a/test/registered/kernels/benchmark/quantization/bench_awq_dequantize.py +++ b/test/registered/kernels/benchmark/quantization/bench_awq_dequantize.py @@ -15,13 +15,6 @@ register_cuda_ci( est_time=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" ) -try: - from sgl_kernel import awq_dequantize as aot_awq_dequantize - - AOT_AVAILABLE = True -except ImportError: - AOT_AVAILABLE = False - IS_CI = is_in_ci() if IS_CI: @@ -34,59 +27,16 @@ else: configs = list(itertools.product(qweight_row_range, qweight_cols_range)) -def check_correctness(): - if not AOT_AVAILABLE: - print("sgl_kernel AOT not available, skipping correctness check") - return - - qweight_row, qweight_col = 128, 16 - device = torch.device("cuda") - qweight = torch.randint( - 0, - torch.iinfo(torch.int32).max, - (qweight_row, qweight_col), - dtype=torch.int32, - device=device, - ) - group_size = qweight_row - scales_row = qweight_row // group_size - scales_col = qweight_col * 8 - scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device) - qzeros = torch.randint( - 0, - torch.iinfo(torch.int32).max, - (scales_row, qweight_col), - dtype=torch.int32, - device=device, - ) - - jit_out = jit_awq_dequantize(qweight, scales, qzeros) - aot_out = aot_awq_dequantize(qweight, scales, qzeros) - torch.cuda.synchronize() - torch.testing.assert_close(jit_out, aot_out, rtol=0, atol=0) - print("Correctness check passed (JIT vs AOT)") - - -if AOT_AVAILABLE: - line_vals = ["jit", "aot"] - line_names = ["JIT Kernel", "AOT Kernel"] - styles = [("blue", "-"), ("green", "-")] -else: - line_vals = ["jit"] - line_names = ["JIT Kernel"] - styles = [("blue", "-")] - - @triton.testing.perf_report( triton.testing.Benchmark( x_names=["qweight_row", "qweight_col"], x_vals=configs, line_arg="provider", - line_vals=line_vals, - line_names=line_names, - styles=styles, + line_vals=["jit"], + line_names=["JIT Kernel"], + styles=[("blue", "-")], ylabel="us", - plot_name="awq-dequantize-jit-vs-aot", + plot_name="awq-dequantize-jit", args={}, ) ) @@ -111,16 +61,12 @@ def benchmark(qweight_row, qweight_col, provider): device=device, ) - if provider == "jit": - fn = lambda: jit_awq_dequantize(qweight, scales, qzeros) - elif provider == "aot": - fn = lambda: aot_awq_dequantize(qweight, scales, qzeros) - else: + if provider != "jit": raise ValueError(f"Unknown provider: {provider}") + fn = lambda: jit_awq_dequantize(qweight, scales, qzeros) return run_benchmark(fn) if __name__ == "__main__": - check_correctness() benchmark.run(print_data=True) diff --git a/test/registered/kernels/ops/quantization/test_awq_dequantize.py b/test/registered/kernels/ops/quantization/test_awq_dequantize.py index f0c47c826..3425e7ba1 100644 --- a/test/registered/kernels/ops/quantization/test_awq_dequantize.py +++ b/test/registered/kernels/ops/quantization/test_awq_dequantize.py @@ -11,13 +11,6 @@ from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=9, stage="base-b-kernel-unit", runner_config="1-gpu-large") -try: - from sgl_kernel import awq_dequantize as aot_awq_dequantize - - AOT_AVAILABLE = True -except ImportError: - AOT_AVAILABLE = False - def reverse_awq_order(t: torch.Tensor): bits = 4 @@ -117,54 +110,5 @@ def test_awq_dequantize_jit_vs_torch( ) -@pytest.mark.parametrize( - "qweight_row,qweight_col,is_bf16_act", - list( - itertools.product( - [128, 256, 512, 1024, 3584], - [16, 32, 64, 128, 448], - [True, False], - ) - ), -) -def test_awq_dequantize_jit_vs_aot( - qweight_row: int, qweight_col: int, is_bf16_act: bool -): - if not AOT_AVAILABLE: - pytest.skip("sgl_kernel AOT not available") - - device = torch.device("cuda") - qweight = torch.randint( - 0, - torch.iinfo(torch.int32).max, - (qweight_row, qweight_col), - dtype=torch.int32, - device=device, - ) - group_size = qweight_row - scales_row = qweight_row // group_size - scales_col = qweight_col * 8 - - if is_bf16_act: - scales = torch.rand(scales_row, scales_col, dtype=torch.bfloat16, device=device) - else: - scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device) - - qzeros = torch.randint( - 0, - torch.iinfo(torch.int32).max, - (scales_row, qweight_col), - dtype=torch.int32, - device=device, - ) - - # Run both implementations - aot_out = aot_awq_dequantize(qweight, scales, qzeros) - jit_out = jit_awq_dequantize(qweight, scales, qzeros) - - # Bitwise equality - torch.testing.assert_close(jit_out, aot_out, rtol=0, atol=0) - - if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/test/registered/unit/layers/quantization/test_gptq_scheme_attach.py b/test/registered/unit/layers/quantization/test_gptq_scheme_attach.py deleted file mode 100644 index 1828b29f0..000000000 --- a/test/registered/unit/layers/quantization/test_gptq_scheme_attach.py +++ /dev/null @@ -1,60 +0,0 @@ -"""GPTQ builds its per-layer scheme lazily in `create_weights`, so the layer has -to declare `scheme = None` for the `is None` probe to see it. - -Regression: the probe used to be `hasattr(layer, "scheme")`, which degraded to -always-true once `LinearBase` grew that class default -- the scheme was never -built and every GPTQ model died with ``'NoneType' object has no attribute -'create_weights'``. -""" - -import unittest - -import torch - -from sglang.srt.layers.linear import LinearBase, ReplicatedLinear -from sglang.srt.layers.moe.fused_moe_triton import FusedMoE -from sglang.srt.layers.quantization.gptq.gptq import GPTQConfig -from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding -from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase - -register_cpu_ci(est_time=11, suite="base-a-test-cpu") - -_GPTQ_CHECKPOINT_CONFIG = { - "bits": 4, - "group_size": 128, - "desc_act": False, - "lm_head": False, - "dynamic": {}, - "checkpoint_format": "gptq", - "true_sequential": True, - "static_groups": False, -} - - -class TestGPTQSchemeAttach(CustomTestCase): - def test_linear_layer_gets_a_scheme(self): - layer = ReplicatedLinear( - input_size=256, - output_size=128, - bias=False, - params_dtype=torch.float16, - quant_config=GPTQConfig.from_config(_GPTQ_CHECKPOINT_CONFIG), - prefix="model.layers.0.mlp.down_proj", - ) - self.assertIsNotNone(layer.scheme) - self.assertTrue(hasattr(layer, "qweight")) - - def test_scheme_default_is_declared_on_every_quantizable_layer_base(self): - """`get_linear_quant_method` hands a linear method a `LinearBase` or a - quantized `ParallelLMHead`; `GPTQMarlinConfig` hands - `GPTQMarlinMoEMethod` a bare `FusedMoE`. The MoE attach has no e2e - coverage, so this is its only guard. - """ - self.assertIsNone(LinearBase.scheme) - self.assertIsNone(VocabParallelEmbedding.scheme) - self.assertIsNone(FusedMoE.scheme) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/unit/server_args/test_page_major_backend_allowlist.py b/test/registered/unit/server_args/test_page_major_backend_allowlist.py index a8ac16cd6..5cb295b14 100644 --- a/test/registered/unit/server_args/test_page_major_backend_allowlist.py +++ b/test/registered/unit/server_args/test_page_major_backend_allowlist.py @@ -101,7 +101,7 @@ class TestPageMajorBackendAllowlist(unittest.TestCase): # MLA-family kernels that must never leak into the MHA arm. MLA_ONLY_BACKENDS = ("trtllm_mla", "cutedsl_mla", "tokenspeed_mla", "flashmla") # No kernel-facing-id wiring anywhere: must stay rejected until they get one. - UNWIRED_BACKENDS = ("cutlass_mla", "aiter") + UNWIRED_BACKENDS = ("aiter",) def test_triton_allowed_on_every_arm(self): """Triton reads both view families, so it is the one backend neither diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index fff1cdc6b..0387dfb6b 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -1423,7 +1423,6 @@ class TestGoldenModelOverrides(_IsolatedPublish): from sglang.srt.arg_groups.overrides import ( ResolvedView, _attention_backend_default, - _attention_backend_dual_chunk, _attention_backend_fa3_fp8_fallback, _attention_backend_platform_fallbacks, ) @@ -1459,19 +1458,6 @@ class TestGoldenModelOverrides(_IsolatedPublish): with override_platform(has_amx=True): self.assertEqual(_attention_backend_platform_fallbacks(view), {}) - # dual-chunk config: mismatched explicit backend raises verbatim - def _mc(dual): - return SimpleNamespace( - _model_config=SimpleNamespace( - hf_config=SimpleNamespace(dual_chunk_attention_config=dual) - ), - attention_backend="fa3", - ) - - with self.assertRaises(ValueError): - _attention_backend_dual_chunk(ResolvedView(_mc({"a": 1}))) - self.assertEqual(_attention_backend_dual_chunk(ResolvedView(_mc(None))), {}) - def test_dllm_platform_paths_at_callable_level(self): from sglang.srt.arg_groups.overrides import ( ResolvedView, @@ -2668,16 +2654,6 @@ class TestGoldenModelOverrides(_IsolatedPublish): ), {"page_size": 64}, ) - # chained: cutlass_mla decode -> 128, then trtllm_mha prefill keeps 128 - self.assertEqual( - _mla_backend_page_constraints( - _view( - decode_attention_backend="cutlass_mla", - prefill_attention_backend="trtllm_mha", - ) - ), - {"page_size": 128}, - ) # no matching backend: nothing declared self.assertEqual(_mla_backend_page_constraints(_view()), {})