Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Khoa Pham
co-authored by
Claude Opus 5
Claude Opus 5
parent
1109e44305
commit
c66a285c94
@@ -0,0 +1,412 @@
|
||||
// Radix top-k core adapted from https://github.com/tile-ai/tilelang/blob/main/examples/deepseek_v32/topk_selector.py.
|
||||
// This JIT variant adds pool expansion and page-table/ragged-offset transforms.
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
#include <bit>
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <cuda_fp16.h>
|
||||
|
||||
namespace sglang {
|
||||
namespace {
|
||||
|
||||
#ifndef C10_LIKELY
|
||||
#define C10_LIKELY(expr) (__builtin_expect(static_cast<bool>(expr), 1))
|
||||
#endif
|
||||
|
||||
#ifndef SGL_GROUP_TOPK
|
||||
#define SGL_GROUP_TOPK 256
|
||||
#endif
|
||||
|
||||
inline constexpr int kGroupTopK = SGL_GROUP_TOPK;
|
||||
inline constexpr int kThreadsPerBlock = 1024;
|
||||
|
||||
inline constexpr std::size_t kSmem = 8 * 1024 * sizeof(uint32_t); // 32KB (bytes)
|
||||
|
||||
struct FastTopKParams {
|
||||
const float* __restrict__ input; // [B, input_stride]
|
||||
const int32_t* __restrict__ row_starts; // [B] or nullptr
|
||||
int32_t* __restrict__ indices; // unused here (kept for layout parity)
|
||||
const int32_t* __restrict__ lengths; // [B]
|
||||
int64_t input_stride;
|
||||
};
|
||||
|
||||
__device__ __forceinline__ auto convert_to_uint8(float x) -> uint8_t {
|
||||
__half h = __float2half_rn(x);
|
||||
uint16_t bits = __half_as_ushort(h);
|
||||
uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits) : static_cast<uint16_t>(bits | 0x8000);
|
||||
return static_cast<uint8_t>(key >> 8);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ auto convert_to_uint32(float x) -> uint32_t {
|
||||
uint32_t bits = __float_as_uint(x);
|
||||
return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
|
||||
}
|
||||
|
||||
template <int K>
|
||||
__device__ void
|
||||
fast_topk_cuda_tl_impl(const float* __restrict__ input, int* __restrict__ index, int row_start, int length) {
|
||||
// We assume length > K here, or it will crash
|
||||
int topk = K;
|
||||
constexpr auto BLOCK_SIZE = 1024;
|
||||
constexpr auto RADIX = 256;
|
||||
constexpr auto SMEM_INPUT_SIZE = kSmem / (2 * sizeof(int));
|
||||
|
||||
alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128];
|
||||
alignas(128) __shared__ int s_counter;
|
||||
alignas(128) __shared__ int s_threshold_bin_id;
|
||||
alignas(128) __shared__ int s_num_input[2];
|
||||
|
||||
auto& s_histogram = s_histogram_buf[0];
|
||||
extern __shared__ int s_input_idx[][SMEM_INPUT_SIZE];
|
||||
|
||||
const int tx = threadIdx.x;
|
||||
|
||||
if (tx < RADIX + 1) s_histogram[tx] = 0;
|
||||
__syncthreads();
|
||||
|
||||
for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
|
||||
const auto bin = convert_to_uint8(input[idx + row_start]);
|
||||
::atomicAdd(&s_histogram[bin], 1);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto run_cumsum = [&] {
|
||||
#pragma unroll 8
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
static_assert(1 << 8 == RADIX);
|
||||
if (C10_LIKELY(tx < RADIX)) {
|
||||
const auto j = 1 << i;
|
||||
const auto k = i & 1;
|
||||
auto value = s_histogram_buf[k][tx];
|
||||
if (tx < RADIX - j) {
|
||||
value += s_histogram_buf[k][tx + j];
|
||||
}
|
||||
s_histogram_buf[k ^ 1][tx] = value;
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
};
|
||||
|
||||
run_cumsum();
|
||||
if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) {
|
||||
s_threshold_bin_id = tx;
|
||||
s_num_input[0] = 0;
|
||||
s_counter = 0;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto threshold_bin = s_threshold_bin_id;
|
||||
topk -= s_histogram[threshold_bin + 1];
|
||||
|
||||
if (topk == 0) {
|
||||
for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
|
||||
const auto bin = static_cast<int>(convert_to_uint8(input[idx + row_start]));
|
||||
if (bin > threshold_bin) {
|
||||
const auto pos = ::atomicAdd(&s_counter, 1);
|
||||
index[pos] = idx;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
return;
|
||||
} else {
|
||||
__syncthreads();
|
||||
if (tx < RADIX + 1) {
|
||||
s_histogram[tx] = 0;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
|
||||
const auto raw_input = input[idx + row_start];
|
||||
const auto bin = static_cast<int>(convert_to_uint8(raw_input));
|
||||
if (bin > threshold_bin) {
|
||||
const auto pos = ::atomicAdd(&s_counter, 1);
|
||||
index[pos] = idx;
|
||||
} else if (bin == threshold_bin) {
|
||||
const auto pos = ::atomicAdd(&s_num_input[0], 1);
|
||||
if (C10_LIKELY(pos < SMEM_INPUT_SIZE)) {
|
||||
s_input_idx[0][pos] = idx;
|
||||
const auto bin = convert_to_uint32(raw_input);
|
||||
const auto sub_bin = (bin >> 24) & 0xFF;
|
||||
::atomicAdd(&s_histogram[sub_bin], 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
#pragma unroll 4
|
||||
for (int round = 0; round < 4; ++round) {
|
||||
__shared__ int s_last_remain;
|
||||
const auto r_idx = round % 2;
|
||||
|
||||
const auto _raw_num_input = s_num_input[r_idx];
|
||||
const auto num_input = (_raw_num_input < int(SMEM_INPUT_SIZE)) ? _raw_num_input : int(SMEM_INPUT_SIZE);
|
||||
|
||||
run_cumsum();
|
||||
if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) {
|
||||
s_threshold_bin_id = tx;
|
||||
s_num_input[r_idx ^ 1] = 0;
|
||||
s_last_remain = topk - s_histogram[tx + 1];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto threshold_bin = s_threshold_bin_id;
|
||||
topk -= s_histogram[threshold_bin + 1];
|
||||
|
||||
if (topk == 0) {
|
||||
for (int i = tx; i < num_input; i += BLOCK_SIZE) {
|
||||
const auto idx = s_input_idx[r_idx][i];
|
||||
const auto offset = 24 - round * 8;
|
||||
const auto bin = (convert_to_uint32(input[idx + row_start]) >> offset) & 0xFF;
|
||||
if (bin > threshold_bin) {
|
||||
const auto pos = ::atomicAdd(&s_counter, 1);
|
||||
index[pos] = idx;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
break;
|
||||
} else {
|
||||
__syncthreads();
|
||||
if (tx < RADIX + 1) {
|
||||
s_histogram[tx] = 0;
|
||||
}
|
||||
__syncthreads();
|
||||
for (int i = tx; i < num_input; i += BLOCK_SIZE) {
|
||||
const auto idx = s_input_idx[r_idx][i];
|
||||
const auto raw_input = input[idx + row_start];
|
||||
const auto offset = 24 - round * 8;
|
||||
const auto bin = (convert_to_uint32(raw_input) >> offset) & 0xFF;
|
||||
if (bin > threshold_bin) {
|
||||
const auto pos = ::atomicAdd(&s_counter, 1);
|
||||
index[pos] = idx;
|
||||
} else if (bin == threshold_bin) {
|
||||
if (round == 3) {
|
||||
const auto pos = ::atomicAdd(&s_last_remain, -1);
|
||||
if (pos > 0) {
|
||||
index[K - pos] = idx;
|
||||
}
|
||||
} else {
|
||||
const auto pos = ::atomicAdd(&s_num_input[r_idx ^ 1], 1);
|
||||
if (C10_LIKELY(pos < SMEM_INPUT_SIZE)) {
|
||||
s_input_idx[r_idx ^ 1][pos] = idx;
|
||||
const auto bin = convert_to_uint32(raw_input);
|
||||
const auto sub_bin = (bin >> (offset - 8)) & 0xFF;
|
||||
::atomicAdd(&s_histogram[sub_bin], 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__device__ __forceinline__ int32_t transform_kpool_token(
|
||||
int32_t raw_token,
|
||||
const int32_t* __restrict__ page_table_entry,
|
||||
const int32_t* __restrict__ topk_indices_offset,
|
||||
int32_t offset) {
|
||||
if (page_table_entry != nullptr) {
|
||||
return page_table_entry[raw_token];
|
||||
}
|
||||
if (topk_indices_offset != nullptr) {
|
||||
return raw_token + offset;
|
||||
}
|
||||
return raw_token;
|
||||
}
|
||||
|
||||
template <int K>
|
||||
__global__ __launch_bounds__(kThreadsPerBlock) void kpool_topk_transform_kernel(
|
||||
const __grid_constant__ FastTopKParams params,
|
||||
int32_t* __restrict__ dst_token_indices,
|
||||
const int64_t dst_stride,
|
||||
const int32_t pool_size,
|
||||
const int32_t token_topk,
|
||||
const int32_t out_cols,
|
||||
const int32_t* __restrict__ page_table,
|
||||
const int64_t page_table_stride,
|
||||
const int32_t* __restrict__ page_table_row_index,
|
||||
const int32_t* __restrict__ topk_indices_offset,
|
||||
const int32_t* __restrict__ seq_lens) {
|
||||
const auto& [input, row_starts, _, lengths, input_stride] = params;
|
||||
const auto bid = static_cast<uint64_t>(blockIdx.x);
|
||||
const auto tid = threadIdx.x;
|
||||
const auto row_start = row_starts == nullptr ? 0 : row_starts[bid];
|
||||
const auto length = lengths[bid];
|
||||
const auto score = input + bid * input_stride;
|
||||
const auto dst = dst_token_indices + bid * dst_stride;
|
||||
const auto page_table_row = page_table_row_index == nullptr ? bid : static_cast<uint64_t>(page_table_row_index[bid]);
|
||||
const auto page_table_entry = page_table == nullptr ? nullptr : page_table + page_table_row * page_table_stride;
|
||||
const auto offset = topk_indices_offset == nullptr ? 0 : topk_indices_offset[bid];
|
||||
const bool append_tail = seq_lens != nullptr;
|
||||
const auto full_pool_token_len = length * pool_size;
|
||||
const auto history_len = full_pool_token_len < token_topk ? full_pool_token_len : token_topk;
|
||||
const auto tail_count = append_tail ? seq_lens[bid] % pool_size : 0;
|
||||
|
||||
if (length <= K) {
|
||||
for (int col = tid; col < out_cols; col += kThreadsPerBlock) {
|
||||
if (col < history_len) {
|
||||
const auto group_rank = col / pool_size;
|
||||
const auto slot = col % pool_size;
|
||||
const auto raw_token = group_rank * pool_size + slot;
|
||||
dst[col] = transform_kpool_token(raw_token, page_table_entry, topk_indices_offset, offset);
|
||||
} else if (append_tail && col < history_len + tail_count) {
|
||||
const auto raw_token = length * pool_size + (col - history_len);
|
||||
dst[col] = transform_kpool_token(raw_token, page_table_entry, topk_indices_offset, offset);
|
||||
} else {
|
||||
dst[col] = -1;
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
__shared__ int s_indices[K];
|
||||
fast_topk_cuda_tl_impl<K>(score, s_indices, row_start, length);
|
||||
for (int col = tid; col < out_cols; col += kThreadsPerBlock) {
|
||||
if (col < history_len) {
|
||||
const auto group_rank = col / pool_size;
|
||||
const auto group_id = s_indices[group_rank];
|
||||
const auto slot = col % pool_size;
|
||||
const auto raw_token = group_id * pool_size + slot;
|
||||
dst[col] = transform_kpool_token(raw_token, page_table_entry, topk_indices_offset, offset);
|
||||
} else if (append_tail && col < history_len + tail_count) {
|
||||
const auto raw_token = length * pool_size + (col - history_len);
|
||||
dst[col] = transform_kpool_token(raw_token, page_table_entry, topk_indices_offset, offset);
|
||||
} else {
|
||||
dst[col] = -1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <auto* f, std::size_t kMaxDynamicSMEM>
|
||||
void setup_kernel_smem_once(host::DebugInfo where = {}) {
|
||||
[[maybe_unused]]
|
||||
static const auto result = [] {
|
||||
const auto fptr = std::bit_cast<const void*>(f);
|
||||
return ::cudaFuncSetAttribute(fptr, ::cudaFuncAttributeMaxDynamicSharedMemorySize, kMaxDynamicSMEM);
|
||||
}();
|
||||
host::RuntimeDeviceCheck(result, where);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
const T* optional_data_ptr(const tvm::ffi::Optional<tvm::ffi::TensorView>& opt) {
|
||||
if (!opt.has_value()) {
|
||||
return nullptr;
|
||||
}
|
||||
return static_cast<const T*>(opt.value().data_ptr());
|
||||
}
|
||||
|
||||
struct KpoolTopKTransformKernel {
|
||||
static constexpr auto kernel = kpool_topk_transform_kernel<kGroupTopK>;
|
||||
|
||||
static void transform(
|
||||
const tvm::ffi::TensorView score,
|
||||
const tvm::ffi::TensorView lengths,
|
||||
const tvm::ffi::TensorView dst_token_indices,
|
||||
const int64_t pool_size,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> page_table_opt,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> topk_indices_offset_opt,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> row_starts_opt,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> seq_lens_opt,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> page_table_row_index_opt) {
|
||||
using namespace host;
|
||||
|
||||
auto B = SymbolicSize{"batch_size"};
|
||||
auto S = SymbolicSize{"score_stride"};
|
||||
auto out_cols_sym = SymbolicSize{"out_cols"};
|
||||
auto device = SymbolicDevice{};
|
||||
device.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({B, -1}).with_strides({S, 1}).with_dtype<float>().with_device(device).verify(score);
|
||||
TensorMatcher({B}).with_dtype<int32_t>().with_device(device).verify(lengths);
|
||||
TensorMatcher({B, out_cols_sym}).with_dtype<int32_t>().with_device(device).verify(dst_token_indices);
|
||||
|
||||
RuntimeCheck(pool_size > 1, "pool_size must be > 1, got ", pool_size);
|
||||
RuntimeCheck(
|
||||
!(page_table_opt.has_value() && topk_indices_offset_opt.has_value()),
|
||||
"page_table and topk_indices_offset are mutually exclusive");
|
||||
RuntimeCheck(
|
||||
!page_table_row_index_opt.has_value() || page_table_opt.has_value(),
|
||||
"page_table_row_index requires page_table");
|
||||
|
||||
const auto out_cols = static_cast<int32_t>(out_cols_sym.unwrap());
|
||||
const auto tail_cols = seq_lens_opt.has_value() ? static_cast<int32_t>(pool_size) - 1 : 0;
|
||||
RuntimeCheck(out_cols > tail_cols, "dst_token_indices columns ", out_cols, " must exceed tail ", tail_cols);
|
||||
const auto token_topk = out_cols - tail_cols;
|
||||
RuntimeCheck(token_topk % static_cast<int32_t>(pool_size) == 0, "token_topk must be a multiple of pool_size");
|
||||
RuntimeCheck(
|
||||
token_topk / static_cast<int32_t>(pool_size) == kGroupTopK,
|
||||
"this module is built for group_topk=",
|
||||
kGroupTopK,
|
||||
" but got ",
|
||||
token_topk / static_cast<int32_t>(pool_size));
|
||||
|
||||
const auto batch_size = static_cast<uint32_t>(B.unwrap());
|
||||
|
||||
int64_t page_table_stride = 0;
|
||||
const int32_t* page_table_ptr = nullptr;
|
||||
if (page_table_opt.has_value()) {
|
||||
auto P = SymbolicSize{"page_table_stride"};
|
||||
if (page_table_row_index_opt.has_value()) {
|
||||
auto page_table_rows = SymbolicSize{"page_table_rows"};
|
||||
TensorMatcher({page_table_rows, -1})
|
||||
.with_strides({P, 1})
|
||||
.with_dtype<int32_t>()
|
||||
.with_device(device)
|
||||
.verify(page_table_opt.value());
|
||||
} else {
|
||||
TensorMatcher({B, -1}).with_strides({P, 1}).with_dtype<int32_t>().with_device(device).verify(
|
||||
page_table_opt.value());
|
||||
}
|
||||
page_table_ptr = static_cast<const int32_t*>(page_table_opt.value().data_ptr());
|
||||
page_table_stride = static_cast<int64_t>(P.unwrap());
|
||||
}
|
||||
|
||||
if (topk_indices_offset_opt.has_value()) {
|
||||
TensorMatcher({B}).with_dtype<int32_t>().with_device(device).verify(topk_indices_offset_opt.value());
|
||||
}
|
||||
if (row_starts_opt.has_value()) {
|
||||
TensorMatcher({B}).with_dtype<int32_t>().with_device(device).verify(row_starts_opt.value());
|
||||
}
|
||||
if (seq_lens_opt.has_value()) {
|
||||
TensorMatcher({B}).with_dtype<int32_t>().with_device(device).verify(seq_lens_opt.value());
|
||||
}
|
||||
if (page_table_row_index_opt.has_value()) {
|
||||
TensorMatcher({B}).with_dtype<int32_t>().with_device(device).verify(page_table_row_index_opt.value());
|
||||
}
|
||||
|
||||
const auto params = FastTopKParams{
|
||||
.input = static_cast<const float*>(score.data_ptr()),
|
||||
.row_starts = optional_data_ptr<int32_t>(row_starts_opt),
|
||||
.indices = nullptr,
|
||||
.lengths = static_cast<const int32_t*>(lengths.data_ptr()),
|
||||
.input_stride = static_cast<int64_t>(S.unwrap()),
|
||||
};
|
||||
|
||||
setup_kernel_smem_once<kernel, kSmem>();
|
||||
LaunchKernel(batch_size, kThreadsPerBlock, device.unwrap(), kSmem)(
|
||||
kernel,
|
||||
params,
|
||||
static_cast<int32_t*>(dst_token_indices.data_ptr()),
|
||||
static_cast<int64_t>(dst_token_indices.strides()[0]),
|
||||
static_cast<int32_t>(pool_size),
|
||||
token_topk,
|
||||
out_cols,
|
||||
page_table_ptr,
|
||||
page_table_stride,
|
||||
optional_data_ptr<int32_t>(page_table_row_index_opt),
|
||||
optional_data_ptr<int32_t>(topk_indices_offset_opt),
|
||||
optional_data_ptr<int32_t>(seq_lens_opt));
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
} // namespace sglang
|
||||
@@ -95,6 +95,7 @@ for _mod, _fn in [
|
||||
("dsa.triton_sparse_mla", "triton_sparse_mla_fwd"),
|
||||
("dsa.transform_index", "transform_index_page_table_prefill"),
|
||||
("dsa.transform_index", "transform_index_page_table_decode"),
|
||||
("dsa.transform_index", "prepare_trtllm_nope_sparse_metadata"),
|
||||
("dsa.cp_split", "dsa_cp_round_robin_split_q_seqs_kernel"),
|
||||
("dsv4.fp4_indexer", "quantize_fp4_indexer_tensor"),
|
||||
("dsv4.fp4_indexer", "store_fp4_index_k_cache"),
|
||||
|
||||
@@ -14,6 +14,58 @@ def transform_index_page_table_decode(**kwargs):
|
||||
return transform_index_page_table_decode_fast(**kwargs)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def prepare_trtllm_nope_sparse_metadata_kernel(
|
||||
page_table_ptr: torch.Tensor,
|
||||
topk_lens_ptr: torch.Tensor,
|
||||
row_stride: tl.constexpr,
|
||||
TOPK: tl.constexpr,
|
||||
BLOCK_TOPK: tl.constexpr,
|
||||
):
|
||||
row = tl.program_id(0)
|
||||
offsets = tl.arange(0, BLOCK_TOPK)
|
||||
valid = offsets < TOPK
|
||||
indices = tl.load(
|
||||
page_table_ptr + row * row_stride + offsets,
|
||||
mask=valid,
|
||||
other=-1,
|
||||
)
|
||||
topk_len = tl.sum((valid & (indices >= 0)).to(tl.int32), axis=0)
|
||||
|
||||
# TRTLLM-GEN's native H512 dynamic sparse kernel produces NaNs for an
|
||||
# empty row. CUDA-graph padding rows are never consumed, so point them at
|
||||
# a valid dummy token and run a one-element attention instead.
|
||||
is_empty = topk_len == 0
|
||||
tl.store(page_table_ptr + row * row_stride, 0, mask=is_empty)
|
||||
tl.store(topk_lens_ptr + row, tl.maximum(topk_len, 1))
|
||||
|
||||
|
||||
def prepare_trtllm_nope_sparse_metadata(
|
||||
page_table: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Build per-query active top-k lengths for native H512 TRTLLM-GEN MLA.
|
||||
|
||||
``page_table`` must contain packed valid token locations followed by ``-1``
|
||||
padding. The tensor is only modified for fully empty CUDA-graph padding
|
||||
rows, whose first entry is replaced with the valid dummy token location 0.
|
||||
"""
|
||||
assert page_table.ndim == 2
|
||||
assert page_table.dtype == torch.int32
|
||||
assert page_table.is_contiguous()
|
||||
num_rows, topk = page_table.shape
|
||||
topk_lens = torch.empty(num_rows, dtype=torch.int32, device=page_table.device)
|
||||
block_topk = triton.next_power_of_2(topk)
|
||||
prepare_trtllm_nope_sparse_metadata_kernel[(num_rows,)](
|
||||
page_table,
|
||||
topk_lens,
|
||||
page_table.stride(0),
|
||||
TOPK=topk,
|
||||
BLOCK_TOPK=block_topk,
|
||||
num_warps=8,
|
||||
)
|
||||
return topk_lens
|
||||
|
||||
|
||||
def _allocate_prefill_result(
|
||||
topk_indices: torch.Tensor,
|
||||
real_num_tokens: int,
|
||||
|
||||
@@ -1155,6 +1155,7 @@ def chunk_kda_fwd(
|
||||
cu_seqlens=cu_seqlens,
|
||||
chunk_size=chunk_size,
|
||||
chunk_indices=chunk_indices,
|
||||
safe_gate=lower_bound is not None,
|
||||
fuse_diagonal=_small_grid,
|
||||
fuse_recompute=_small_grid,
|
||||
)
|
||||
@@ -1210,6 +1211,7 @@ def chunk_kda(
|
||||
dt_bias: Optional[torch.Tensor] = None,
|
||||
lower_bound: Optional[float] = None,
|
||||
output_intermediate_states: bool = False,
|
||||
beta_is_raw: bool = False,
|
||||
**kwargs,
|
||||
):
|
||||
if scale is None:
|
||||
@@ -1219,6 +1221,9 @@ def chunk_kda(
|
||||
q = l2norm_fwd(q.contiguous())
|
||||
k = l2norm_fwd(k.contiguous())
|
||||
|
||||
if beta_is_raw:
|
||||
beta = beta.float().sigmoid()
|
||||
|
||||
# Returns o [B, T, H, V] when output_intermediate_states=False, or (o, h [B, NT, H, V, K]) when output_intermediate_states=True.
|
||||
return chunk_kda_fwd(
|
||||
q=q,
|
||||
|
||||
@@ -1314,6 +1314,7 @@ def chunk_kda(
|
||||
dt_bias: torch.Tensor | None = None,
|
||||
lower_bound: float | None = None,
|
||||
output_intermediate_states: bool = False,
|
||||
beta_is_raw: bool = False,
|
||||
**kwargs: object,
|
||||
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Match the public forward contract of SGLang's Triton ``chunk_kda``."""
|
||||
@@ -1327,6 +1328,8 @@ def chunk_kda(
|
||||
raise ValueError("g and beta must cover every q token")
|
||||
g = g[:, :num_tokens]
|
||||
beta = beta[:, :num_tokens]
|
||||
if beta_is_raw:
|
||||
beta = beta.float().sigmoid()
|
||||
if num_tokens == 1:
|
||||
# Tracing constant-folds size-one dimensions, but the resulting kernel
|
||||
# can share a cache entry with longer inputs. Keep T=1 on Triton so a
|
||||
|
||||
@@ -180,6 +180,27 @@ def mla_quantize_and_rope_for_fp8(
|
||||
return q_out, k_nope_out, k_rope_out
|
||||
|
||||
|
||||
def mla_quantize_for_fp8_no_rope(
|
||||
q_nope: torch.Tensor,
|
||||
q_rope: torch.Tensor,
|
||||
k_nope: torch.Tensor,
|
||||
k_rope: torch.Tensor,
|
||||
kv_lora_rank: int,
|
||||
qk_rope_head_dim: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
attn_dtype = torch.float8_e4m3fn
|
||||
q_len, num_heads = q_rope.shape[:2]
|
||||
q_out = q_rope.new_empty(
|
||||
q_len,
|
||||
num_heads,
|
||||
kv_lora_rank + qk_rope_head_dim,
|
||||
dtype=attn_dtype,
|
||||
)
|
||||
q_out[..., :kv_lora_rank] = q_nope.to(attn_dtype)
|
||||
q_out[..., kv_lora_rank:] = q_rope.to(attn_dtype)
|
||||
return q_out, k_nope.to(attn_dtype), k_rope.to(attn_dtype)
|
||||
|
||||
|
||||
def mla_quantize_without_rope_for_fp8(
|
||||
q_nope: torch.Tensor,
|
||||
q_rope: torch.Tensor,
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
@@ -81,6 +83,40 @@ def set_mla_kv_buffer_kernel(
|
||||
tl.extra.cuda.gdc_launch_dependents()
|
||||
|
||||
|
||||
@triton.jit
|
||||
def set_mla_kv_buffer_kernel_norope(
|
||||
kv_buffer_ptr,
|
||||
cache_k_nope_ptr,
|
||||
loc_ptr,
|
||||
buffer_stride: tl.constexpr,
|
||||
nope_stride: tl.constexpr,
|
||||
nope_dim: tl.constexpr,
|
||||
BLOCK: tl.constexpr,
|
||||
USE_GDC: tl.constexpr = False,
|
||||
):
|
||||
pid_loc = tl.program_id(0)
|
||||
pid_blk = tl.program_id(1)
|
||||
|
||||
base = pid_blk * BLOCK
|
||||
offs = base + tl.arange(0, BLOCK)
|
||||
mask = offs < nope_dim
|
||||
|
||||
if USE_GDC:
|
||||
tl.extra.cuda.gdc_wait()
|
||||
|
||||
loc = tl.load(loc_ptr + pid_loc).to(tl.int64)
|
||||
dst_ptr = kv_buffer_ptr + loc * buffer_stride + offs
|
||||
|
||||
src = tl.load(
|
||||
cache_k_nope_ptr + pid_loc * nope_stride + offs,
|
||||
mask=mask,
|
||||
)
|
||||
tl.store(dst_ptr, src, mask=mask)
|
||||
|
||||
if USE_GDC:
|
||||
tl.extra.cuda.gdc_launch_dependents()
|
||||
|
||||
|
||||
# Above this loc count the TMA bulk-store path overtakes the single-CTA-per-loc
|
||||
# Triton kernel. Below it, Triton with BLOCK = next_pow2(total_dim) (one CTA
|
||||
# does the whole row in one tile, no boundary fan-out) is the winning fallback.
|
||||
@@ -92,7 +128,7 @@ def _set_mla_kv_buffer_impl(
|
||||
kv_buffer: torch.Tensor,
|
||||
loc: torch.Tensor,
|
||||
cache_k_nope: torch.Tensor,
|
||||
cache_k_rope: torch.Tensor,
|
||||
cache_k_rope: Optional[torch.Tensor] = None,
|
||||
*,
|
||||
reserved_skip_index: int,
|
||||
dcp_world_size: int,
|
||||
@@ -127,6 +163,28 @@ def _set_mla_kv_buffer_impl(
|
||||
Shared body of the two entry points below; the owner rule reaches it as
|
||||
``1, 0`` (nothing to select) or as the live topology.
|
||||
"""
|
||||
has_rope = cache_k_rope is not None and cache_k_rope.numel() > 0
|
||||
n_loc = loc.numel()
|
||||
nope_dim = cache_k_nope.shape[-1]
|
||||
|
||||
if not has_rope:
|
||||
BLOCK = triton.next_power_of_2(nope_dim)
|
||||
grid = (n_loc, 1)
|
||||
pdl_kwargs = (
|
||||
{"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
|
||||
)
|
||||
set_mla_kv_buffer_kernel_norope[grid](
|
||||
kv_buffer,
|
||||
cache_k_nope,
|
||||
loc,
|
||||
kv_buffer.stride(0),
|
||||
cache_k_nope.stride(0),
|
||||
nope_dim,
|
||||
BLOCK=BLOCK,
|
||||
**pdl_kwargs,
|
||||
)
|
||||
return
|
||||
|
||||
from sglang.kernels.ops.kvcache.set_mla_kv_buffer import (
|
||||
can_use_set_mla_kv_buffer,
|
||||
)
|
||||
@@ -134,7 +192,6 @@ def _set_mla_kv_buffer_impl(
|
||||
set_mla_kv_buffer as jit_set_mla_kv_buffer,
|
||||
)
|
||||
|
||||
n_loc = loc.numel()
|
||||
nope_bytes = cache_k_nope.shape[-1] * cache_k_nope.element_size()
|
||||
rope_bytes = cache_k_rope.shape[-1] * cache_k_rope.element_size()
|
||||
if (
|
||||
@@ -157,7 +214,6 @@ def _set_mla_kv_buffer_impl(
|
||||
# ``set_mla_kv_buffer_kernel`` handles the over-allocation past total_dim
|
||||
# via the offs<total_dim mask). Beats BLOCK=128 by 60-2700 ns across the
|
||||
# 2 <= bs <= 512 range on GB300.
|
||||
nope_dim = cache_k_nope.shape[-1]
|
||||
rope_dim = cache_k_rope.shape[-1]
|
||||
total_dim = nope_dim + rope_dim
|
||||
BLOCK = triton.next_power_of_2(total_dim)
|
||||
@@ -437,18 +493,51 @@ def get_mla_kv_buffer_kernel(
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def get_mla_kv_buffer_kernel_norope(
|
||||
kv_buffer_ptr,
|
||||
cache_k_nope_ptr,
|
||||
loc_ptr,
|
||||
buffer_stride: tl.constexpr,
|
||||
nope_stride: tl.constexpr,
|
||||
nope_dim: tl.constexpr,
|
||||
):
|
||||
pid_loc = tl.program_id(0)
|
||||
loc = tl.load(loc_ptr + pid_loc).to(tl.int64)
|
||||
loc_src_ptr = kv_buffer_ptr + loc * buffer_stride
|
||||
|
||||
nope_offs = tl.arange(0, nope_dim)
|
||||
nope_src = tl.load(loc_src_ptr + nope_offs)
|
||||
tl.store(
|
||||
cache_k_nope_ptr + pid_loc * nope_stride + nope_offs,
|
||||
nope_src,
|
||||
)
|
||||
|
||||
|
||||
def get_mla_kv_buffer_triton(
|
||||
kv_buffer: torch.Tensor,
|
||||
loc: torch.Tensor,
|
||||
cache_k_nope: torch.Tensor,
|
||||
cache_k_rope: torch.Tensor,
|
||||
cache_k_rope: Optional[torch.Tensor] = None,
|
||||
):
|
||||
# The source data type will be implicitly converted to the target data type.
|
||||
nope_dim = cache_k_nope.shape[-1] # 512
|
||||
rope_dim = cache_k_rope.shape[-1] # 64
|
||||
n_loc = loc.numel()
|
||||
grid = (n_loc,)
|
||||
|
||||
has_rope = cache_k_rope is not None and cache_k_rope.numel() > 0
|
||||
if not has_rope:
|
||||
get_mla_kv_buffer_kernel_norope[grid](
|
||||
kv_buffer,
|
||||
cache_k_nope,
|
||||
loc,
|
||||
kv_buffer.stride(0),
|
||||
cache_k_nope.stride(0),
|
||||
nope_dim,
|
||||
)
|
||||
return
|
||||
|
||||
rope_dim = cache_k_rope.shape[-1] # 64
|
||||
get_mla_kv_buffer_kernel[grid](
|
||||
kv_buffer,
|
||||
cache_k_nope,
|
||||
|
||||
@@ -1734,6 +1734,190 @@ def mhc_fused_post_pre(
|
||||
)
|
||||
|
||||
|
||||
def hc_expand(x: torch.Tensor, n: int) -> torch.Tensor:
|
||||
return x.repeat(1, n)
|
||||
|
||||
|
||||
def hc_contract(x: torch.Tensor, n: int) -> torch.Tensor:
|
||||
return x.unflatten(-1, (n, -1)).mean(dim=-2)
|
||||
|
||||
|
||||
def _mhc_pre_torch(
|
||||
residual: torch.Tensor,
|
||||
fn: torch.Tensor,
|
||||
hc_scale: torch.Tensor,
|
||||
hc_base: torch.Tensor,
|
||||
rms_eps: float,
|
||||
hc_pre_eps: float,
|
||||
hc_sinkhorn_eps: float,
|
||||
hc_post_mult_value: float,
|
||||
sinkhorn_repeat: int,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
import torch.nn.functional as F
|
||||
|
||||
s, n, h = residual.shape
|
||||
dtype = residual.dtype
|
||||
|
||||
x_flat = residual.view(s, n * h).float()
|
||||
rsqrt = torch.rsqrt(x_flat.square().mean(-1, keepdim=True) + rms_eps)
|
||||
mixes = F.linear(x_flat, fn) * rsqrt
|
||||
|
||||
pre_raw = mixes[:, :n]
|
||||
post_raw = mixes[:, n : 2 * n]
|
||||
comb_raw = mixes[:, 2 * n :].view(s, n, n)
|
||||
pre_base = hc_base[:n]
|
||||
post_base = hc_base[n : 2 * n]
|
||||
comb_base = hc_base[2 * n :].view(n, n)
|
||||
|
||||
pre = torch.sigmoid(pre_raw * hc_scale[0] + pre_base) + hc_pre_eps
|
||||
post = hc_post_mult_value * torch.sigmoid(post_raw * hc_scale[1] + post_base)
|
||||
comb = comb_raw * hc_scale[2] + comb_base
|
||||
|
||||
comb = comb.softmax(-1) + hc_sinkhorn_eps
|
||||
comb = comb / (comb.sum(-2, keepdim=True) + hc_sinkhorn_eps)
|
||||
for _ in range(sinkhorn_repeat - 1):
|
||||
comb = comb / (comb.sum(-1, keepdim=True) + hc_sinkhorn_eps)
|
||||
comb = comb / (comb.sum(-2, keepdim=True) + hc_sinkhorn_eps)
|
||||
|
||||
layer_input = (pre.unsqueeze(-1) * residual.float()).sum(dim=1).to(dtype)
|
||||
return post.unsqueeze(-1), comb, layer_input
|
||||
|
||||
|
||||
def _mhc_post_torch(
|
||||
x: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
post_layer_mix: torch.Tensor,
|
||||
comb_res_mix: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
out = post_layer_mix * x.unsqueeze(1) + (
|
||||
comb_res_mix.unsqueeze(-1) * residual.unsqueeze(2)
|
||||
).sum(dim=1)
|
||||
return out.type_as(x)
|
||||
|
||||
|
||||
@torch._dynamo.disable
|
||||
def _mhc_pre_dispatch(
|
||||
residual: torch.Tensor,
|
||||
fn: torch.Tensor,
|
||||
hc_scale: torch.Tensor,
|
||||
hc_base: torch.Tensor,
|
||||
rms_eps: float,
|
||||
hc_pre_eps: float,
|
||||
hc_sinkhorn_eps: float,
|
||||
hc_post_mult_value: float,
|
||||
sinkhorn_repeat: int,
|
||||
norm_weight: torch.Tensor | None = None,
|
||||
norm_eps: float | None = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, bool]:
|
||||
assert residual.dim() == 3, f"residual must be (s, n, h); got {residual.shape}"
|
||||
if not envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get():
|
||||
post_mix, comb_mix, layer_input = _mhc_pre_torch(
|
||||
residual=residual,
|
||||
fn=fn,
|
||||
hc_scale=hc_scale,
|
||||
hc_base=hc_base,
|
||||
rms_eps=rms_eps,
|
||||
hc_pre_eps=hc_pre_eps,
|
||||
hc_sinkhorn_eps=hc_sinkhorn_eps,
|
||||
hc_post_mult_value=hc_post_mult_value,
|
||||
sinkhorn_repeat=sinkhorn_repeat,
|
||||
)
|
||||
return post_mix, comb_mix, layer_input, False
|
||||
|
||||
post_mix, comb_mix, layer_input = mhc_pre(
|
||||
residual=residual,
|
||||
fn=fn,
|
||||
hc_scale=hc_scale,
|
||||
hc_base=hc_base,
|
||||
rms_eps=rms_eps,
|
||||
hc_pre_eps=hc_pre_eps,
|
||||
hc_sinkhorn_eps=hc_sinkhorn_eps,
|
||||
hc_post_mult_value=hc_post_mult_value,
|
||||
sinkhorn_repeat=sinkhorn_repeat,
|
||||
norm_weight=norm_weight,
|
||||
norm_eps=norm_eps,
|
||||
)
|
||||
return post_mix, comb_mix, layer_input, norm_weight is not None
|
||||
|
||||
|
||||
@torch._dynamo.disable
|
||||
def _mhc_post_dispatch(
|
||||
x: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
post_layer_mix: torch.Tensor,
|
||||
comb_res_mix: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
assert x.dim() == 2 and residual.dim() == 3
|
||||
assert post_layer_mix.dim() == 3 and comb_res_mix.dim() == 3
|
||||
if not envs.SGLANG_OPT_USE_TILELANG_MHC_POST.get():
|
||||
return _mhc_post_torch(x, residual, post_layer_mix, comb_res_mix)
|
||||
return mhc_post(x, residual, post_layer_mix, comb_res_mix)
|
||||
|
||||
|
||||
def hc_pre(
|
||||
x: torch.Tensor,
|
||||
hc_fn: torch.Tensor,
|
||||
hc_scale: torch.Tensor,
|
||||
hc_base: torch.Tensor,
|
||||
hc_mult: int,
|
||||
rms_eps: float,
|
||||
hc_eps: float,
|
||||
sinkhorn_iters: int,
|
||||
post_mult_value: float = 2.0,
|
||||
hc_norm_weight: torch.Tensor | None = None,
|
||||
out_norm_weight: torch.Tensor | None = None,
|
||||
out_norm_eps: float | None = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, bool]:
|
||||
s, total = x.shape
|
||||
hidden_size = total // hc_mult
|
||||
if x.numel() == 0:
|
||||
empty_layer_input = x.new_zeros((s, hidden_size))
|
||||
empty_h_res = torch.zeros(
|
||||
(s, hc_mult * hc_mult), device=x.device, dtype=torch.float32
|
||||
)
|
||||
empty_h_post = torch.zeros((s, hc_mult), device=x.device, dtype=torch.float32)
|
||||
return empty_layer_input, empty_h_res, empty_h_post, False
|
||||
|
||||
fn = hc_fn if hc_norm_weight is None else hc_fn * hc_norm_weight
|
||||
residual_3d = x.view(s, hc_mult, hidden_size)
|
||||
post_mix, comb_mix, layer_input, norm_fused = _mhc_pre_dispatch(
|
||||
residual=residual_3d,
|
||||
fn=fn,
|
||||
hc_scale=hc_scale,
|
||||
hc_base=hc_base,
|
||||
rms_eps=rms_eps,
|
||||
hc_pre_eps=hc_eps,
|
||||
hc_sinkhorn_eps=hc_eps,
|
||||
hc_post_mult_value=post_mult_value,
|
||||
sinkhorn_repeat=sinkhorn_iters,
|
||||
norm_weight=out_norm_weight,
|
||||
norm_eps=out_norm_eps,
|
||||
)
|
||||
return (
|
||||
layer_input,
|
||||
comb_mix.reshape(s, hc_mult * hc_mult),
|
||||
post_mix.reshape(s, hc_mult),
|
||||
norm_fused,
|
||||
)
|
||||
|
||||
|
||||
def hc_post(
|
||||
x: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
h_post: torch.Tensor,
|
||||
h_res: torch.Tensor,
|
||||
hc_mult: int,
|
||||
) -> torch.Tensor:
|
||||
s, hidden_size = x.shape
|
||||
if s == 0:
|
||||
return x.new_zeros((s, hc_mult * hidden_size))
|
||||
residual = residual.view(s, hc_mult, hidden_size)
|
||||
h_post = h_post.view(s, hc_mult, 1)
|
||||
h_res = h_res.view(s, hc_mult, hc_mult)
|
||||
out = _mhc_post_dispatch(x, residual, h_post, h_res)
|
||||
return out.view(s, -1)
|
||||
|
||||
|
||||
def npu_hc_pre(
|
||||
x: torch.Tensor,
|
||||
hc_fn: torch.Tensor,
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.utils import cache_once, load_jit
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
SUPPORTED_GROUP_TOPK = (128, 160, 192, 224, 256, 512)
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_kpool_topk_transform_module(group_topk: int) -> Module:
|
||||
assert group_topk in SUPPORTED_GROUP_TOPK, (
|
||||
"fast_kpool_topk_transform supports pool-level topk "
|
||||
f"{SUPPORTED_GROUP_TOPK}, got {group_topk}"
|
||||
)
|
||||
return load_jit(
|
||||
f"kpool_topk_transform_{group_topk}",
|
||||
cuda_files=["dsa/kpool_topk_transform.cuh"],
|
||||
cuda_wrappers=[("kpool_topk_transform", "KpoolTopKTransformKernel::transform")],
|
||||
extra_cuda_cflags=[f"-DSGL_GROUP_TOPK={group_topk}"],
|
||||
)
|
||||
|
||||
|
||||
def fast_kpool_topk_transform_fused(
|
||||
score: torch.Tensor,
|
||||
lengths: torch.Tensor,
|
||||
pool_size: int,
|
||||
topk: int,
|
||||
page_table: Optional[torch.Tensor] = None,
|
||||
topk_indices_offset: Optional[torch.Tensor] = None,
|
||||
row_starts: Optional[torch.Tensor] = None,
|
||||
seq_lens: Optional[torch.Tensor] = None,
|
||||
page_table_row_index: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
assert topk % pool_size == 0
|
||||
group_topk = topk // pool_size
|
||||
assert group_topk in SUPPORTED_GROUP_TOPK, (
|
||||
"fast_kpool_topk_transform supports pool-level topk "
|
||||
f"{SUPPORTED_GROUP_TOPK}, got {group_topk}"
|
||||
)
|
||||
assert score.dim() == 2
|
||||
assert page_table is None or topk_indices_offset is None
|
||||
assert page_table_row_index is None or page_table is not None
|
||||
if seq_lens is not None:
|
||||
assert seq_lens.dim() == 1
|
||||
assert seq_lens.shape[0] == score.shape[0]
|
||||
if page_table_row_index is not None:
|
||||
assert page_table_row_index.dim() == 1
|
||||
assert page_table_row_index.shape[0] == score.shape[0]
|
||||
|
||||
out_cols = topk + (pool_size - 1 if seq_lens is not None else 0)
|
||||
dst_token_indices = score.new_empty((score.shape[0], out_cols), dtype=torch.int32)
|
||||
|
||||
module = _jit_kpool_topk_transform_module(group_topk)
|
||||
module.kpool_topk_transform(
|
||||
score,
|
||||
lengths,
|
||||
dst_token_indices,
|
||||
pool_size,
|
||||
page_table,
|
||||
topk_indices_offset,
|
||||
row_starts,
|
||||
seq_lens,
|
||||
page_table_row_index,
|
||||
)
|
||||
return dst_token_indices
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,852 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, List, NamedTuple, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.attention.dsa.kpool_fp8_index import (
|
||||
INDEX_HEAD_DIM,
|
||||
build_pooled_page_table_64,
|
||||
kpool_build_ragged_layout,
|
||||
kpool_max_closed_pools,
|
||||
update_kpool_write_plan_cuda_graph,
|
||||
)
|
||||
from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp
|
||||
from sglang.srt.model_executor.forward_context import get_req_to_token_pool
|
||||
from sglang.srt.runtime_context import get_parallel
|
||||
from sglang.srt.utils import is_cuda
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.attention.dsa.dsa_topk_backend import TopkTransformMethod
|
||||
from sglang.srt.layers.attention.dsa_backend import DSAMetadata
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
|
||||
|
||||
|
||||
_RAGGED_SCRATCH_K_U8: Optional[torch.Tensor] = None
|
||||
_RAGGED_SCRATCH_K_SCALE: Optional[torch.Tensor] = None
|
||||
|
||||
|
||||
def _get_ragged_scratch(
|
||||
total_k_rows: int, device: torch.device
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
global _RAGGED_SCRATCH_K_U8, _RAGGED_SCRATCH_K_SCALE
|
||||
cur = _RAGGED_SCRATCH_K_U8
|
||||
grow = (
|
||||
cur is None
|
||||
or cur.device.type != device.type
|
||||
or (device.index is not None and cur.device.index != device.index)
|
||||
or cur.shape[0] < total_k_rows
|
||||
)
|
||||
if grow:
|
||||
_RAGGED_SCRATCH_K_U8 = torch.empty(
|
||||
(total_k_rows, INDEX_HEAD_DIM), dtype=torch.uint8, device=device
|
||||
)
|
||||
_RAGGED_SCRATCH_K_SCALE = torch.empty(
|
||||
(total_k_rows,), dtype=torch.float32, device=device
|
||||
)
|
||||
return _RAGGED_SCRATCH_K_U8[:total_k_rows], _RAGGED_SCRATCH_K_SCALE[:total_k_rows]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PoolWriteRows:
|
||||
req: torch.Tensor
|
||||
pool_id: torch.Tensor
|
||||
n_from_tail: torch.Tensor
|
||||
chunk_src: torch.Tensor
|
||||
tail_logical_base: torch.Tensor
|
||||
write_loc: torch.Tensor
|
||||
|
||||
@property
|
||||
def is_empty(self) -> bool:
|
||||
return self.pool_id.shape[0] == 0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TailWriteRows:
|
||||
req: torch.Tensor
|
||||
dst_logical_start: torch.Tensor
|
||||
chunk_src: torch.Tensor
|
||||
n_write: torch.Tensor
|
||||
|
||||
@property
|
||||
def is_empty(self) -> bool:
|
||||
return self.req.shape[0] == 0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class KPoolCpInfo:
|
||||
size: int
|
||||
rank: int
|
||||
owner_rank: torch.Tensor
|
||||
local_write_mask: torch.Tensor
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class KPoolExtendPlan:
|
||||
writes: PoolWriteRows
|
||||
tails: TailWriteRows
|
||||
pooled_seq_lens_expanded: torch.Tensor
|
||||
seq_lens_expanded: torch.Tensor
|
||||
ragged_concat_page_table: torch.Tensor
|
||||
ragged_q_ks: torch.Tensor
|
||||
ragged_q_ke: torch.Tensor
|
||||
ragged_total_k_rows: int
|
||||
ragged_k_u8: Optional[torch.Tensor]
|
||||
ragged_k_scale: Optional[torch.Tensor]
|
||||
ragged_paged_page_table: Optional[torch.Tensor]
|
||||
ragged_paged_page_table_row_index: Optional[torch.Tensor]
|
||||
cp: Optional[KPoolCpInfo] = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class KPoolWritePlan:
|
||||
"""``write_loc[b, p]`` is the compression destination for candidate closed
|
||||
pool ``base_pool[b] + p``; the kernel decides which candidates closed."""
|
||||
|
||||
req: torch.Tensor
|
||||
write_start: torch.Tensor
|
||||
tail_logical_start: torch.Tensor
|
||||
write_loc: torch.Tensor # int64 [B, max_closed_pools]
|
||||
num_draft_tokens: int
|
||||
pool_seqlens_per_q: Optional[torch.Tensor] = None
|
||||
seqlens_per_q: Optional[torch.Tensor] = None
|
||||
pool_schedule_metadata: Optional[torch.Tensor] = None
|
||||
effective_n_per_batch: Optional[torch.Tensor] = None
|
||||
|
||||
|
||||
def _is_kpool_layout_enabled(pool_size: int, real_page_size: int) -> bool:
|
||||
return pool_size > 1 and real_page_size == 64 and real_page_size % pool_size == 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class _KPoolCpuPlan:
|
||||
pool_batch_idx: List[int] = field(default_factory=list)
|
||||
pool_req: List[int] = field(default_factory=list)
|
||||
pool_pool_id: List[int] = field(default_factory=list)
|
||||
pool_n_from_tail: List[int] = field(default_factory=list)
|
||||
pool_chunk_src: List[int] = field(default_factory=list)
|
||||
pool_tail_logical_base: List[int] = field(default_factory=list)
|
||||
|
||||
tail_req: List[int] = field(default_factory=list)
|
||||
tail_dst_logical_start: List[int] = field(default_factory=list)
|
||||
tail_chunk_src: List[int] = field(default_factory=list)
|
||||
tail_n_write: List[int] = field(default_factory=list)
|
||||
|
||||
ragged_q_len: List[int] = field(default_factory=list)
|
||||
ragged_pool_pages: List[int] = field(default_factory=list)
|
||||
cu_pages_excl: List[int] = field(default_factory=list)
|
||||
cu_q_len_excl: List[int] = field(default_factory=list)
|
||||
total_pool_pages: int = 0
|
||||
|
||||
|
||||
class _KPoolDecompose(NamedTuple):
|
||||
first_slot: int
|
||||
base_pool: int
|
||||
n_pool: int
|
||||
tail_n_write: int
|
||||
|
||||
|
||||
def _decompose_compress(start: int, length: int, pool_size: int) -> _KPoolDecompose:
|
||||
first_slot = start % pool_size
|
||||
base_pool = start // pool_size
|
||||
n_pool = (start + length) // pool_size - base_pool
|
||||
consumed = max(0, n_pool * pool_size - first_slot)
|
||||
tail_n = length - consumed
|
||||
return _KPoolDecompose(
|
||||
first_slot=first_slot,
|
||||
base_pool=base_pool,
|
||||
n_pool=n_pool,
|
||||
tail_n_write=tail_n,
|
||||
)
|
||||
|
||||
|
||||
def _append_compress_rows(
|
||||
plan: _KPoolCpuPlan,
|
||||
pool_size: int,
|
||||
batch_size: int,
|
||||
extend_seq_lens_cpu: List[int],
|
||||
seq_lens_cpu: List[int],
|
||||
req_pool_indices_cpu: List[int],
|
||||
) -> None:
|
||||
q_offset = 0
|
||||
for i in range(batch_size):
|
||||
q_len = extend_seq_lens_cpu[i]
|
||||
assert q_len > 0, f"extend_seq_lens_cpu[{i}] = {q_len}; expected > 0"
|
||||
|
||||
seq_len = seq_lens_cpu[i]
|
||||
req = req_pool_indices_cpu[i]
|
||||
d = _decompose_compress(seq_len - q_len, q_len, pool_size)
|
||||
|
||||
if d.n_pool > 0:
|
||||
plan.pool_batch_idx.extend([i] * d.n_pool)
|
||||
plan.pool_req.extend([req] * d.n_pool)
|
||||
plan.pool_pool_id.extend(range(d.base_pool, d.base_pool + d.n_pool))
|
||||
plan.pool_n_from_tail.append(d.first_slot)
|
||||
plan.pool_n_from_tail.extend([0] * (d.n_pool - 1))
|
||||
bulk_start = q_offset + pool_size - d.first_slot
|
||||
plan.pool_chunk_src.append(q_offset)
|
||||
plan.pool_chunk_src.extend(
|
||||
range(bulk_start, bulk_start + (d.n_pool - 1) * pool_size, pool_size)
|
||||
)
|
||||
plan.pool_tail_logical_base.extend(
|
||||
range(
|
||||
d.base_pool * pool_size,
|
||||
(d.base_pool + d.n_pool) * pool_size,
|
||||
pool_size,
|
||||
)
|
||||
)
|
||||
|
||||
if d.tail_n_write > 0:
|
||||
consumed = q_len - d.tail_n_write
|
||||
plan.tail_req.append(req)
|
||||
plan.tail_dst_logical_start.append(seq_len - q_len + consumed)
|
||||
plan.tail_chunk_src.append(q_offset + consumed)
|
||||
plan.tail_n_write.append(d.tail_n_write)
|
||||
|
||||
q_offset += q_len
|
||||
|
||||
|
||||
def _append_local_rows(
|
||||
plan: _KPoolCpuPlan,
|
||||
pool_size: int,
|
||||
slots_per_page: int,
|
||||
local_extend_seq_lens_cpu: List[int],
|
||||
local_seq_lens_cpu: List[int],
|
||||
) -> None:
|
||||
q_offset = 0
|
||||
for q_len, seq_len in zip(
|
||||
local_extend_seq_lens_cpu, local_seq_lens_cpu, strict=True
|
||||
):
|
||||
assert q_len > 0, f"local_extend_seq_lens_cpu has non-positive {q_len = }"
|
||||
plan.ragged_q_len.append(q_len)
|
||||
pool_seq_len = seq_len // pool_size
|
||||
pool_pages_i = (pool_seq_len + slots_per_page - 1) // slots_per_page
|
||||
plan.ragged_pool_pages.append(pool_pages_i)
|
||||
plan.cu_pages_excl.append(plan.total_pool_pages)
|
||||
plan.cu_q_len_excl.append(q_offset)
|
||||
plan.total_pool_pages += pool_pages_i
|
||||
q_offset += q_len
|
||||
|
||||
|
||||
def _kpool_cpu_plan(
|
||||
forward_batch: ForwardBatch,
|
||||
pool_size: int,
|
||||
slots_per_page: int,
|
||||
*,
|
||||
local_extend_seq_lens_cpu: Optional[List[int]] = None,
|
||||
local_seq_lens_cpu: Optional[List[int]] = None,
|
||||
) -> _KPoolCpuPlan:
|
||||
plan = _KPoolCpuPlan()
|
||||
|
||||
extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu
|
||||
if isinstance(extend_seq_lens_cpu, torch.Tensor):
|
||||
extend_seq_lens_cpu = extend_seq_lens_cpu.tolist()
|
||||
seq_lens_cpu = forward_batch.seq_lens_cpu.tolist()
|
||||
req_pool_indices_cpu = forward_batch.req_pool_indices.tolist()
|
||||
|
||||
_append_compress_rows(
|
||||
plan,
|
||||
pool_size,
|
||||
forward_batch.batch_size,
|
||||
extend_seq_lens_cpu,
|
||||
seq_lens_cpu,
|
||||
req_pool_indices_cpu,
|
||||
)
|
||||
|
||||
if local_extend_seq_lens_cpu is None:
|
||||
local_extend_seq_lens_cpu = extend_seq_lens_cpu
|
||||
local_seq_lens_cpu = seq_lens_cpu
|
||||
|
||||
_append_local_rows(
|
||||
plan,
|
||||
pool_size,
|
||||
slots_per_page,
|
||||
local_extend_seq_lens_cpu,
|
||||
local_seq_lens_cpu,
|
||||
)
|
||||
return plan
|
||||
|
||||
|
||||
def _kpool_plan_to_gpu(
|
||||
cpu: _KPoolCpuPlan,
|
||||
forward_batch: ForwardBatch,
|
||||
full_real_page_table: torch.Tensor,
|
||||
local_real_page_table: torch.Tensor,
|
||||
local_seqlens_expanded: torch.Tensor,
|
||||
local_req_pool_indices: torch.Tensor,
|
||||
pool_size: int,
|
||||
slots_per_page: int,
|
||||
topk_transform_method: TopkTransformMethod,
|
||||
) -> KPoolExtendPlan:
|
||||
from sglang.srt.layers.attention.dsa.dsa_topk_backend import TopkTransformMethod
|
||||
|
||||
device = forward_batch.seq_lens.device
|
||||
n_pool = len(cpu.pool_pool_id)
|
||||
n_tail = len(cpu.tail_req)
|
||||
n_rag = len(cpu.ragged_q_len)
|
||||
|
||||
total_pool_pages = cpu.total_pool_pages
|
||||
ragged_total_k_rows = total_pool_pages * slots_per_page
|
||||
|
||||
need_paged = (
|
||||
topk_transform_method == TopkTransformMethod.PAGED
|
||||
and envs.SGLANG_DSA_FUSE_TOPK.get()
|
||||
and n_rag > 0
|
||||
)
|
||||
|
||||
i64_total = 4 * n_pool + 2 * n_tail
|
||||
if i64_total > 0:
|
||||
i64_cpu = torch.tensor(
|
||||
cpu.pool_req
|
||||
+ cpu.pool_pool_id
|
||||
+ cpu.pool_chunk_src
|
||||
+ cpu.pool_batch_idx
|
||||
+ cpu.tail_req
|
||||
+ cpu.tail_chunk_src,
|
||||
dtype=torch.int64,
|
||||
pin_memory=True,
|
||||
)
|
||||
i64_gpu = i64_cpu.to(device, non_blocking=True)
|
||||
c = 0
|
||||
pool_req_t = i64_gpu[c : c + n_pool]
|
||||
c += n_pool
|
||||
pool_pool_id_t = i64_gpu[c : c + n_pool]
|
||||
c += n_pool
|
||||
pool_chunk_src_t = i64_gpu[c : c + n_pool]
|
||||
c += n_pool
|
||||
pool_batch_idx_t = i64_gpu[c : c + n_pool]
|
||||
c += n_pool
|
||||
tail_req_t = i64_gpu[c : c + n_tail]
|
||||
c += n_tail
|
||||
tail_chunk_src_t = i64_gpu[c : c + n_tail]
|
||||
else:
|
||||
empty_i64 = torch.empty((0,), dtype=torch.int64, device=device)
|
||||
pool_req_t = pool_pool_id_t = pool_chunk_src_t = pool_batch_idx_t = empty_i64
|
||||
tail_req_t = tail_chunk_src_t = empty_i64
|
||||
|
||||
i32_total = 2 * n_pool + 2 * n_tail + 4 * n_rag
|
||||
if i32_total > 0:
|
||||
i32_cpu = torch.tensor(
|
||||
cpu.pool_n_from_tail
|
||||
+ cpu.pool_tail_logical_base
|
||||
+ cpu.tail_dst_logical_start
|
||||
+ cpu.tail_n_write
|
||||
+ cpu.ragged_pool_pages
|
||||
+ cpu.ragged_q_len
|
||||
+ cpu.cu_pages_excl
|
||||
+ cpu.cu_q_len_excl,
|
||||
dtype=torch.int32,
|
||||
pin_memory=True,
|
||||
)
|
||||
i32_gpu = i32_cpu.to(device, non_blocking=True)
|
||||
c = 0
|
||||
pool_n_from_tail_t = i32_gpu[c : c + n_pool]
|
||||
c += n_pool
|
||||
pool_tail_logical_base_t = i32_gpu[c : c + n_pool]
|
||||
c += n_pool
|
||||
tail_dst_logical_start_t = i32_gpu[c : c + n_tail]
|
||||
c += n_tail
|
||||
tail_n_write_t = i32_gpu[c : c + n_tail]
|
||||
c += n_tail
|
||||
ragged_pool_pages_t = i32_gpu[c : c + n_rag]
|
||||
c += n_rag
|
||||
ragged_q_len_t = i32_gpu[c : c + n_rag]
|
||||
c += n_rag
|
||||
cu_pages_excl_t = i32_gpu[c : c + n_rag]
|
||||
c += n_rag
|
||||
cu_q_len_excl_t = i32_gpu[c : c + n_rag]
|
||||
else:
|
||||
empty_i32 = torch.empty((0,), dtype=torch.int32, device=device)
|
||||
pool_n_from_tail_t = pool_tail_logical_base_t = empty_i32
|
||||
tail_dst_logical_start_t = tail_n_write_t = empty_i32
|
||||
ragged_pool_pages_t = ragged_q_len_t = empty_i32
|
||||
cu_pages_excl_t = cu_q_len_excl_t = empty_i32
|
||||
|
||||
if n_pool > 0:
|
||||
pool_page_group = torch.div(
|
||||
pool_pool_id_t, slots_per_page, rounding_mode="floor"
|
||||
)
|
||||
token_page_row = pool_page_group * pool_size
|
||||
packed_page = full_real_page_table[pool_batch_idx_t, token_page_row].to(
|
||||
torch.int64
|
||||
)
|
||||
pool_write_locs = packed_page * slots_per_page + torch.remainder(
|
||||
pool_pool_id_t, slots_per_page
|
||||
)
|
||||
else:
|
||||
pool_write_locs = torch.empty((0,), dtype=torch.int64, device=device)
|
||||
|
||||
pooled_seq_lens_expanded = torch.div(
|
||||
local_seqlens_expanded, pool_size, rounding_mode="floor"
|
||||
).to(torch.int32)
|
||||
|
||||
if n_rag > 0:
|
||||
(
|
||||
ragged_concat_page_table,
|
||||
ragged_q_ks,
|
||||
ragged_q_ke,
|
||||
) = kpool_build_ragged_layout(
|
||||
full_page_table=local_real_page_table,
|
||||
cu_pages_excl=cu_pages_excl_t,
|
||||
ragged_pool_pages=ragged_pool_pages_t,
|
||||
cu_q_len_excl=cu_q_len_excl_t,
|
||||
ragged_q_len=ragged_q_len_t,
|
||||
pooled_seq_lens_expanded=pooled_seq_lens_expanded,
|
||||
slots_per_page=slots_per_page,
|
||||
total_pool_pages=total_pool_pages,
|
||||
total_q=pooled_seq_lens_expanded.shape[0],
|
||||
pool_size=pool_size,
|
||||
)
|
||||
else:
|
||||
empty_i32_dev = torch.empty((0,), dtype=torch.int32, device=device)
|
||||
ragged_concat_page_table = empty_i32_dev
|
||||
ragged_q_ks = empty_i32_dev
|
||||
ragged_q_ke = empty_i32_dev
|
||||
|
||||
ragged_paged_page_table = None
|
||||
ragged_paged_page_table_row_index = None
|
||||
if need_paged:
|
||||
req_to_token = get_req_to_token_pool().req_to_token
|
||||
ragged_paged_page_table_row_index = torch.repeat_interleave(
|
||||
local_req_pool_indices.to(torch.int32), ragged_q_len_t
|
||||
)
|
||||
ragged_paged_page_table = req_to_token
|
||||
|
||||
if ragged_total_k_rows > 0:
|
||||
ragged_k_u8, ragged_k_scale = _get_ragged_scratch(ragged_total_k_rows, device)
|
||||
else:
|
||||
ragged_k_u8 = None
|
||||
ragged_k_scale = None
|
||||
|
||||
return KPoolExtendPlan(
|
||||
writes=PoolWriteRows(
|
||||
req=pool_req_t,
|
||||
pool_id=pool_pool_id_t,
|
||||
n_from_tail=pool_n_from_tail_t,
|
||||
chunk_src=pool_chunk_src_t,
|
||||
tail_logical_base=pool_tail_logical_base_t,
|
||||
write_loc=pool_write_locs,
|
||||
),
|
||||
tails=TailWriteRows(
|
||||
req=tail_req_t,
|
||||
dst_logical_start=tail_dst_logical_start_t,
|
||||
chunk_src=tail_chunk_src_t,
|
||||
n_write=tail_n_write_t,
|
||||
),
|
||||
pooled_seq_lens_expanded=pooled_seq_lens_expanded,
|
||||
seq_lens_expanded=local_seqlens_expanded,
|
||||
ragged_concat_page_table=ragged_concat_page_table,
|
||||
ragged_q_ks=ragged_q_ks,
|
||||
ragged_q_ke=ragged_q_ke,
|
||||
ragged_total_k_rows=ragged_total_k_rows,
|
||||
ragged_k_u8=ragged_k_u8,
|
||||
ragged_k_scale=ragged_k_scale,
|
||||
ragged_paged_page_table=ragged_paged_page_table,
|
||||
ragged_paged_page_table_row_index=ragged_paged_page_table_row_index,
|
||||
cp=_kpool_cp_owner_rank(forward_batch, n_pool, device),
|
||||
)
|
||||
|
||||
|
||||
def _kpool_cp_owner_rank(
|
||||
forward_batch: ForwardBatch,
|
||||
n_pool: int,
|
||||
device: torch.device,
|
||||
) -> Optional[KPoolCpInfo]:
|
||||
if not dsa_use_prefill_cp(forward_batch):
|
||||
return None
|
||||
|
||||
cp_size = get_parallel().attn_cp_size
|
||||
if cp_size <= 1:
|
||||
return None
|
||||
|
||||
cp_rank = get_parallel().attn_cp_rank
|
||||
if n_pool > 0:
|
||||
owner = torch.arange(n_pool, dtype=torch.int32, device=device) % cp_size
|
||||
local_write_mask = owner == cp_rank
|
||||
else:
|
||||
owner = torch.empty((0,), dtype=torch.int32, device=device)
|
||||
local_write_mask = torch.empty((0,), dtype=torch.bool, device=device)
|
||||
return KPoolCpInfo(
|
||||
size=cp_size,
|
||||
rank=cp_rank,
|
||||
owner_rank=owner,
|
||||
local_write_mask=local_write_mask,
|
||||
)
|
||||
|
||||
|
||||
def init_kpool_extend_metadata(
|
||||
metadata: DSAMetadata,
|
||||
forward_batch: ForwardBatch,
|
||||
*,
|
||||
pool_size: int,
|
||||
real_page_size: int,
|
||||
slots_per_page: int,
|
||||
topk_transform_method: TopkTransformMethod,
|
||||
full_real_page_table: torch.Tensor,
|
||||
full_seqlens_expanded: torch.Tensor,
|
||||
local_real_page_table: Optional[torch.Tensor] = None,
|
||||
local_seqlens_expanded: Optional[torch.Tensor] = None,
|
||||
local_extend_seq_lens_cpu: Optional[List[int]] = None,
|
||||
local_seq_lens_cpu: Optional[List[int]] = None,
|
||||
local_req_pool_indices: Optional[torch.Tensor] = None,
|
||||
) -> DSAMetadata:
|
||||
mode = forward_batch.forward_mode
|
||||
is_extend_like = mode.is_extend_without_speculative() or mode.is_draft_extend_v2()
|
||||
if (
|
||||
not _is_kpool_layout_enabled(pool_size, real_page_size)
|
||||
or not is_extend_like
|
||||
or forward_batch.extend_seq_lens_cpu is None
|
||||
or forward_batch.seq_lens_cpu is None
|
||||
):
|
||||
return metadata
|
||||
|
||||
if local_real_page_table is None:
|
||||
local_real_page_table = full_real_page_table
|
||||
if local_seqlens_expanded is None:
|
||||
local_seqlens_expanded = full_seqlens_expanded
|
||||
if local_req_pool_indices is None:
|
||||
local_req_pool_indices = forward_batch.req_pool_indices
|
||||
|
||||
cpu = _kpool_cpu_plan(
|
||||
forward_batch,
|
||||
pool_size,
|
||||
slots_per_page,
|
||||
local_extend_seq_lens_cpu=local_extend_seq_lens_cpu,
|
||||
local_seq_lens_cpu=local_seq_lens_cpu,
|
||||
)
|
||||
plan = _kpool_plan_to_gpu(
|
||||
cpu,
|
||||
forward_batch,
|
||||
full_real_page_table,
|
||||
local_real_page_table,
|
||||
local_seqlens_expanded,
|
||||
local_req_pool_indices,
|
||||
pool_size,
|
||||
slots_per_page,
|
||||
topk_transform_method,
|
||||
)
|
||||
return dataclasses.replace(metadata, kpool_extend_plan=plan)
|
||||
|
||||
|
||||
_DEEP_GEMM_MODULE = None
|
||||
_DEEP_GEMM_IMPORT_FAILED = False
|
||||
|
||||
|
||||
def _get_deep_gemm():
|
||||
# Cache import failure too, because this helper runs on every graph-replay
|
||||
# metadata refresh.
|
||||
global _DEEP_GEMM_MODULE, _DEEP_GEMM_IMPORT_FAILED
|
||||
if _DEEP_GEMM_MODULE is None and not _DEEP_GEMM_IMPORT_FAILED:
|
||||
try:
|
||||
import deep_gemm
|
||||
except (ImportError, ModuleNotFoundError):
|
||||
_DEEP_GEMM_IMPORT_FAILED = True
|
||||
else:
|
||||
_DEEP_GEMM_MODULE = deep_gemm
|
||||
return _DEEP_GEMM_MODULE
|
||||
|
||||
|
||||
def _compute_pool_schedule_metadata(
|
||||
pool_seqlens: torch.Tensor,
|
||||
*,
|
||||
slots_per_page: int,
|
||||
) -> Optional[torch.Tensor]:
|
||||
if not is_cuda():
|
||||
return None
|
||||
deep_gemm = _get_deep_gemm()
|
||||
if deep_gemm is None:
|
||||
return None
|
||||
return deep_gemm.get_paged_mqa_logits_metadata(
|
||||
pool_seqlens.contiguous().view(-1, 1).clamp(min=1),
|
||||
slots_per_page,
|
||||
deep_gemm.get_num_sms(),
|
||||
)
|
||||
|
||||
|
||||
def init_pooled_paged_mqa_metadata(
|
||||
metadata: DSAMetadata,
|
||||
seqlens_32: torch.Tensor,
|
||||
forward_mode: ForwardMode,
|
||||
*,
|
||||
pool_size: int,
|
||||
real_page_size: int,
|
||||
slots_per_page: int,
|
||||
build_schedule_metadata: bool = True,
|
||||
) -> DSAMetadata:
|
||||
if (
|
||||
not _is_kpool_layout_enabled(pool_size, real_page_size)
|
||||
or not is_cuda()
|
||||
or not forward_mode.is_decode_or_idle()
|
||||
):
|
||||
return metadata
|
||||
|
||||
pool_seqlens = torch.div(seqlens_32, pool_size, rounding_mode="floor").to(
|
||||
torch.int32
|
||||
)
|
||||
pooled_page_table = build_pooled_page_table_64(
|
||||
metadata.real_page_table, pool_size
|
||||
).contiguous()
|
||||
schedule = (
|
||||
_compute_pool_schedule_metadata(
|
||||
pool_seqlens,
|
||||
slots_per_page=slots_per_page,
|
||||
)
|
||||
if build_schedule_metadata
|
||||
else None
|
||||
)
|
||||
return dataclasses.replace(
|
||||
metadata,
|
||||
pooled_index_kpool=pool_size,
|
||||
pooled_cache_seqlens_int32=pool_seqlens,
|
||||
pooled_real_page_table=pooled_page_table,
|
||||
pooled_paged_mqa_schedule_metadata=schedule,
|
||||
)
|
||||
|
||||
|
||||
def update_pooled_paged_mqa_metadata(
|
||||
metadata: DSAMetadata,
|
||||
seqlens_32: torch.Tensor,
|
||||
forward_mode: ForwardMode,
|
||||
*,
|
||||
pool_size: int,
|
||||
real_page_size: int,
|
||||
slots_per_page: int,
|
||||
build_schedule_metadata: bool = True,
|
||||
) -> None:
|
||||
if (
|
||||
not _is_kpool_layout_enabled(pool_size, real_page_size)
|
||||
or not is_cuda()
|
||||
or not forward_mode.is_decode_or_idle()
|
||||
):
|
||||
return
|
||||
|
||||
if (
|
||||
metadata.pooled_index_kpool != pool_size
|
||||
or metadata.pooled_cache_seqlens_int32 is None
|
||||
or metadata.pooled_real_page_table is None
|
||||
):
|
||||
return
|
||||
|
||||
pool_seqlens = torch.div(seqlens_32, pool_size, rounding_mode="floor").to(
|
||||
torch.int32
|
||||
)
|
||||
metadata.pooled_cache_seqlens_int32[: pool_seqlens.shape[0]].copy_(pool_seqlens)
|
||||
pooled_page_table = build_pooled_page_table_64(
|
||||
metadata.real_page_table, pool_size
|
||||
).contiguous()
|
||||
metadata.pooled_real_page_table[
|
||||
: pooled_page_table.shape[0], : pooled_page_table.shape[1]
|
||||
].copy_(pooled_page_table)
|
||||
|
||||
if (
|
||||
build_schedule_metadata
|
||||
and metadata.pooled_paged_mqa_schedule_metadata is not None
|
||||
):
|
||||
new_schedule = _compute_pool_schedule_metadata(
|
||||
metadata.pooled_cache_seqlens_int32,
|
||||
slots_per_page=slots_per_page,
|
||||
)
|
||||
if new_schedule is not None:
|
||||
metadata.pooled_paged_mqa_schedule_metadata.copy_(new_schedule)
|
||||
|
||||
|
||||
def _alloc_kpool_write_plan_buffers(
|
||||
*,
|
||||
max_bs: int,
|
||||
num_draft_tokens: int,
|
||||
pool_size: int,
|
||||
device: torch.device,
|
||||
is_verify: bool,
|
||||
is_v2: bool = False,
|
||||
) -> KPoolWritePlan:
|
||||
max_closed_pools = kpool_max_closed_pools(num_draft_tokens, pool_size)
|
||||
verify_extras = {}
|
||||
if is_verify:
|
||||
n_rows = max_bs * num_draft_tokens
|
||||
verify_extras = dict(
|
||||
pool_seqlens_per_q=torch.zeros(n_rows, dtype=torch.int32, device=device),
|
||||
seqlens_per_q=torch.zeros(n_rows, dtype=torch.int32, device=device),
|
||||
)
|
||||
if is_v2:
|
||||
verify_extras["effective_n_per_batch"] = torch.zeros(
|
||||
max_bs, dtype=torch.int32, device=device
|
||||
)
|
||||
return KPoolWritePlan(
|
||||
req=torch.zeros(max_bs, dtype=torch.int64, device=device),
|
||||
write_start=torch.zeros(max_bs, dtype=torch.int32, device=device),
|
||||
tail_logical_start=torch.zeros(max_bs, dtype=torch.int32, device=device),
|
||||
write_loc=torch.zeros(
|
||||
max_bs, max_closed_pools, dtype=torch.int64, device=device
|
||||
),
|
||||
num_draft_tokens=num_draft_tokens,
|
||||
**verify_extras,
|
||||
)
|
||||
|
||||
|
||||
def init_kpool_write_plan_capture(
|
||||
metadata: DSAMetadata,
|
||||
*,
|
||||
max_bs: int,
|
||||
pool_size: int,
|
||||
real_page_size: int,
|
||||
num_draft_tokens: int,
|
||||
device: torch.device,
|
||||
is_verify: bool,
|
||||
slots_per_page: int,
|
||||
is_v2: bool = False,
|
||||
build_schedule_metadata: bool = True,
|
||||
) -> DSAMetadata:
|
||||
if not _is_kpool_layout_enabled(pool_size, real_page_size) or num_draft_tokens == 0:
|
||||
return metadata
|
||||
|
||||
plan = _alloc_kpool_write_plan_buffers(
|
||||
max_bs=max_bs,
|
||||
num_draft_tokens=num_draft_tokens,
|
||||
pool_size=pool_size,
|
||||
device=device,
|
||||
is_verify=is_verify,
|
||||
is_v2=is_v2,
|
||||
)
|
||||
if is_verify and build_schedule_metadata:
|
||||
schedule = _compute_pool_schedule_metadata(
|
||||
plan.pool_seqlens_per_q,
|
||||
slots_per_page=slots_per_page,
|
||||
)
|
||||
plan = dataclasses.replace(plan, pool_schedule_metadata=schedule)
|
||||
return dataclasses.replace(metadata, kpool_write_plan=plan)
|
||||
|
||||
|
||||
def update_kpool_write_plan(
|
||||
metadata: DSAMetadata,
|
||||
*,
|
||||
write_start: torch.Tensor,
|
||||
req_pool_indices: torch.Tensor,
|
||||
real_page_table: torch.Tensor,
|
||||
pool_size: int,
|
||||
real_page_size: int,
|
||||
num_draft_tokens: int,
|
||||
forward_mode: ForwardMode,
|
||||
slots_per_page: int,
|
||||
effective_n_per_batch: Optional[torch.Tensor] = None,
|
||||
include_deep_gemm_schedule: bool = True,
|
||||
) -> None:
|
||||
if not _is_kpool_layout_enabled(pool_size, real_page_size) or not is_cuda():
|
||||
return
|
||||
is_verify = forward_mode.is_target_verify()
|
||||
is_decode = forward_mode.is_decode_or_idle()
|
||||
is_v2 = forward_mode.is_draft_extend_v2()
|
||||
if not (is_verify or is_decode or is_v2):
|
||||
return
|
||||
|
||||
plan = metadata.kpool_write_plan
|
||||
assert plan is not None, "kpool_write_plan must be allocated before update"
|
||||
update_kpool_write_plan_cuda_graph(
|
||||
write_start=write_start,
|
||||
req_pool_indices=req_pool_indices,
|
||||
real_page_table=real_page_table,
|
||||
req_out=plan.req,
|
||||
write_start_out=plan.write_start,
|
||||
tail_logical_start_out=plan.tail_logical_start,
|
||||
write_loc_out=plan.write_loc,
|
||||
pool_seqlens_per_q_out=plan.pool_seqlens_per_q,
|
||||
seqlens_per_q_out=plan.seqlens_per_q,
|
||||
pool_size=pool_size,
|
||||
num_draft_tokens=num_draft_tokens,
|
||||
slots_per_page=slots_per_page,
|
||||
)
|
||||
|
||||
if (
|
||||
is_v2
|
||||
and effective_n_per_batch is not None
|
||||
and plan.effective_n_per_batch is not None
|
||||
):
|
||||
plan.effective_n_per_batch[: effective_n_per_batch.shape[0]].copy_(
|
||||
effective_n_per_batch.to(torch.int32)
|
||||
)
|
||||
|
||||
# In-graph replay updates plan lengths too late for host schedule construction;
|
||||
# the caller rebuilds the schedule from raw seq_lens out of graph.
|
||||
if include_deep_gemm_schedule and plan.pool_schedule_metadata is not None:
|
||||
new_schedule = _compute_pool_schedule_metadata(
|
||||
plan.pool_seqlens_per_q,
|
||||
slots_per_page=slots_per_page,
|
||||
)
|
||||
if new_schedule is not None:
|
||||
plan.pool_schedule_metadata.copy_(new_schedule)
|
||||
|
||||
|
||||
def refresh_kpool_pool_schedule_from(
|
||||
metadata: DSAMetadata,
|
||||
pool_seqlens_per_q: torch.Tensor,
|
||||
*,
|
||||
slots_per_page: int,
|
||||
) -> None:
|
||||
"""Use an explicit source because the captured plan buffer remains stale
|
||||
until replay."""
|
||||
plan = metadata.kpool_write_plan
|
||||
if plan is None or plan.pool_schedule_metadata is None:
|
||||
return
|
||||
new_schedule = _compute_pool_schedule_metadata(
|
||||
pool_seqlens_per_q,
|
||||
slots_per_page=slots_per_page,
|
||||
)
|
||||
if new_schedule is not None:
|
||||
plan.pool_schedule_metadata.copy_(new_schedule)
|
||||
|
||||
|
||||
def init_kpool_write_plan(
|
||||
metadata: DSAMetadata,
|
||||
forward_batch: ForwardBatch,
|
||||
*,
|
||||
pool_size: int,
|
||||
real_page_size: int,
|
||||
real_page_table: torch.Tensor,
|
||||
num_draft_tokens: int,
|
||||
write_start: torch.Tensor,
|
||||
slots_per_page: int,
|
||||
effective_n_per_batch: Optional[torch.Tensor] = None,
|
||||
build_schedule_metadata: bool = True,
|
||||
) -> DSAMetadata:
|
||||
forward_mode = forward_batch.forward_mode
|
||||
is_verify = forward_mode.is_target_verify()
|
||||
is_decode = forward_mode.is_decode_or_idle()
|
||||
is_v2 = forward_mode.is_draft_extend_v2()
|
||||
is_ring_write = is_verify or is_decode or is_v2
|
||||
if not _is_kpool_layout_enabled(pool_size, real_page_size) or not is_ring_write:
|
||||
return metadata
|
||||
|
||||
pool = getattr(forward_batch, "token_to_kv_pool", None)
|
||||
if (is_verify or is_v2) and pool is not None:
|
||||
assert pool.tail_extra_slots == num_draft_tokens, (
|
||||
f"tail_extra_slots mismatch: pool={pool.tail_extra_slots}, "
|
||||
f"forward={num_draft_tokens}"
|
||||
)
|
||||
|
||||
metadata = init_kpool_write_plan_capture(
|
||||
metadata,
|
||||
max_bs=forward_batch.seq_lens.shape[0],
|
||||
pool_size=pool_size,
|
||||
real_page_size=real_page_size,
|
||||
num_draft_tokens=num_draft_tokens,
|
||||
device=forward_batch.seq_lens.device,
|
||||
is_verify=is_verify or is_v2,
|
||||
slots_per_page=slots_per_page,
|
||||
is_v2=is_v2,
|
||||
build_schedule_metadata=build_schedule_metadata,
|
||||
)
|
||||
update_kpool_write_plan(
|
||||
metadata,
|
||||
write_start=write_start,
|
||||
req_pool_indices=forward_batch.req_pool_indices,
|
||||
real_page_table=real_page_table,
|
||||
pool_size=pool_size,
|
||||
real_page_size=real_page_size,
|
||||
num_draft_tokens=num_draft_tokens,
|
||||
forward_mode=forward_mode,
|
||||
slots_per_page=slots_per_page,
|
||||
effective_n_per_batch=effective_n_per_batch,
|
||||
)
|
||||
return metadata
|
||||
@@ -138,6 +138,8 @@ class CuteDSLKDAKernel(LinearAttnKernelBase):
|
||||
num_tokens = q_n.shape[0]
|
||||
g_in = g[0][:num_tokens] # raw forget gate; activated inside chunk_kda_cutedsl
|
||||
beta_in = beta[0][:num_tokens].to(torch.float32)
|
||||
if kwargs.get("beta_is_raw"):
|
||||
beta_in = beta_in.sigmoid()
|
||||
cu_seqlens = query_start_loc.to(torch.int32)
|
||||
|
||||
# Pool state I/O is fused into the h kernel's TMA load/store: pass the
|
||||
|
||||
@@ -39,6 +39,7 @@ def _triton_fallback(
|
||||
A_log=None,
|
||||
dt_bias=None,
|
||||
lower_bound=None,
|
||||
beta_is_raw=False,
|
||||
return_intermediate_states=False,
|
||||
):
|
||||
"""Fall back to the Triton chunk_kda kernel (handles all preprocessing).
|
||||
@@ -64,6 +65,7 @@ def _triton_fallback(
|
||||
A_log=A_log,
|
||||
dt_bias=dt_bias,
|
||||
lower_bound=lower_bound,
|
||||
beta_is_raw=beta_is_raw,
|
||||
output_intermediate_states=return_intermediate_states,
|
||||
)
|
||||
|
||||
@@ -114,6 +116,7 @@ class FlashKDAKernel(LinearAttnKernelBase):
|
||||
lower_bound: Optional[float] = None,
|
||||
extend_seq_lens_cpu: Optional[list] = None,
|
||||
is_spec_decode: bool = False,
|
||||
beta_is_raw: bool = False,
|
||||
return_intermediate_states: bool = False,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
@@ -136,6 +139,7 @@ class FlashKDAKernel(LinearAttnKernelBase):
|
||||
A_log=A_log,
|
||||
dt_bias=dt_bias,
|
||||
lower_bound=lower_bound,
|
||||
beta_is_raw=beta_is_raw,
|
||||
return_intermediate_states=return_intermediate_states,
|
||||
)
|
||||
|
||||
@@ -152,6 +156,7 @@ class FlashKDAKernel(LinearAttnKernelBase):
|
||||
A_log=A_log,
|
||||
dt_bias=dt_bias,
|
||||
lower_bound=lower_bound,
|
||||
beta_is_raw=beta_is_raw,
|
||||
),
|
||||
None,
|
||||
)
|
||||
@@ -206,6 +211,7 @@ class FlashKDAKernel(LinearAttnKernelBase):
|
||||
A_log: Optional[torch.Tensor] = None,
|
||||
dt_bias: Optional[torch.Tensor] = None,
|
||||
lower_bound: Optional[float] = None,
|
||||
beta_is_raw: bool = False,
|
||||
) -> torch.Tensor:
|
||||
flash_kda = _load_flash_kda()
|
||||
|
||||
@@ -223,12 +229,11 @@ class FlashKDAKernel(LinearAttnKernelBase):
|
||||
v = v.contiguous()
|
||||
g = g.contiguous()
|
||||
|
||||
# KimiDeltaAttention.forward already applies sigmoid to beta on the
|
||||
# prefill path, but flash_kda expects beta LOGITS (it sigmoids
|
||||
# internally). Invert back so the kernel recovers the intended value:
|
||||
# sigmoid(logit(p)) == p. (triton/cuLA consume the post-sigmoid beta.)
|
||||
beta = torch.logit(beta.float().clamp_(1e-7, 1.0 - 1e-7)).to(torch.bfloat16)
|
||||
beta = beta.contiguous()
|
||||
# FlashKDA applies sigmoid internally; invert only the already-activated
|
||||
# Kimi beta path.
|
||||
if not beta_is_raw:
|
||||
beta = torch.logit(beta.float().clamp_(1e-7, 1.0 - 1e-7))
|
||||
beta = beta.to(torch.bfloat16).contiguous()
|
||||
|
||||
# flash_kda wants A_log [H] fp32 and dt_bias [H, K] fp32. The model
|
||||
# stores A_log as [1, 1, H, 1] and dt_bias as 1D [H*K], so reshape both.
|
||||
|
||||
@@ -171,7 +171,7 @@ class TritonKDAKernel(LinearAttnKernelBase):
|
||||
intermediate_states_buffer: torch.Tensor,
|
||||
intermediate_state_indices: torch.Tensor,
|
||||
cache_steps: int,
|
||||
retrieve_parent_token: torch.Tensor,
|
||||
retrieve_parent_token: Optional[torch.Tensor],
|
||||
lower_bound: Optional[float] = None,
|
||||
# fused ReplaySSM ring-write (dense verify only; off elsewhere).
|
||||
cache_ring: bool = False,
|
||||
@@ -229,6 +229,7 @@ class TritonKDAKernel(LinearAttnKernelBase):
|
||||
A_log: Optional[torch.Tensor] = None,
|
||||
dt_bias: Optional[torch.Tensor] = None,
|
||||
lower_bound: Optional[float] = None,
|
||||
beta_is_raw: bool = False,
|
||||
return_intermediate_states: bool = False,
|
||||
**kwargs,
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
@@ -245,5 +246,6 @@ class TritonKDAKernel(LinearAttnKernelBase):
|
||||
A_log=A_log,
|
||||
dt_bias=dt_bias,
|
||||
lower_bound=lower_bound,
|
||||
beta_is_raw=beta_is_raw,
|
||||
output_intermediate_states=return_intermediate_states,
|
||||
)
|
||||
|
||||
@@ -706,6 +706,7 @@ def _compare_prefill(
|
||||
A_log: torch.Tensor | None = None,
|
||||
dt_bias: torch.Tensor | None = None,
|
||||
lower_bound: float | None = None,
|
||||
beta_is_raw: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
batch, tokens, heads, key_dim = q.shape
|
||||
value_dim = v.size(-1)
|
||||
@@ -740,7 +741,8 @@ def _compare_prefill(
|
||||
k_rows = reference_k.view(batch * tokens, heads, key_dim)
|
||||
v_rows = v.view(batch * tokens, heads, value_dim).float()
|
||||
gate_rows = reference_gate.view(batch * tokens, heads, key_dim)
|
||||
beta_rows = beta.view(batch * tokens, heads).float()
|
||||
reference_beta = beta.float().sigmoid() if beta_is_raw else beta
|
||||
beta_rows = reference_beta.view(batch * tokens, heads).float()
|
||||
out_rows = reference_out.view(batch * tokens, heads, value_dim)
|
||||
|
||||
if cu_seqlens is None:
|
||||
@@ -811,6 +813,7 @@ def _compare_prefill(
|
||||
A_log=A_log,
|
||||
dt_bias=dt_bias,
|
||||
lower_bound=lower_bound,
|
||||
beta_is_raw=beta_is_raw,
|
||||
)
|
||||
|
||||
assert helion_out.data_ptr() == helion_v.data_ptr()
|
||||
@@ -852,6 +855,42 @@ def test_fixed_partial_prefill_and_state_pool_contract() -> None:
|
||||
assert torch.equal(helion_state[untouched], state[untouched])
|
||||
|
||||
|
||||
def test_raw_beta_prefill_contract() -> None:
|
||||
torch.manual_seed(811)
|
||||
batch, tokens, heads, key_dim, value_dim = 2, 17, 2, 32, 32
|
||||
q = torch.randn(batch, tokens, heads, key_dim, device="cuda", dtype=torch.bfloat16)
|
||||
# Keep the unnormalized recurrence numerically contractive while still
|
||||
# exercising the no-QK-L2-normalization path. Unit-scale random keys make
|
||||
# (I - beta * k k^T) expansive and obscure the raw-beta contract with
|
||||
# exponentially amplified BF16 round-off.
|
||||
k = torch.randn_like(q) * 0.05
|
||||
v = torch.randn(
|
||||
batch, tokens, heads, value_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
# Keep the recurrent decay contractive so the raw-beta check measures the
|
||||
# sigmoid conversion instead of amplifying BF16 round-off exponentially.
|
||||
gate = -torch.rand_like(q) * 0.2
|
||||
raw_beta = torch.linspace(
|
||||
-2,
|
||||
2,
|
||||
steps=batch * tokens * heads,
|
||||
device="cuda",
|
||||
).reshape(batch, tokens, heads)
|
||||
indices = torch.tensor([3, 1], device="cuda", dtype=torch.int32)
|
||||
state = torch.randn(5, heads, value_dim, key_dim, device="cuda") * 0.01
|
||||
|
||||
_compare_prefill(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
gate,
|
||||
raw_beta,
|
||||
state,
|
||||
indices,
|
||||
beta_is_raw=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("is_varlen", [False, True], ids=["fixed", "varlen"])
|
||||
def test_single_token_prefill_does_not_poison_later_shapes(
|
||||
is_varlen: bool,
|
||||
|
||||
@@ -0,0 +1,213 @@
|
||||
"""CUDA regressions for DSA kpool speculative writes spanning multiple pools."""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.attention.dsa.kpool_fp8_index import (
|
||||
INDEX_HEAD_DIM,
|
||||
kpool_assemble_softmax_rotate_write_cache,
|
||||
kpool_max_closed_pools,
|
||||
kpool_write_tail_and_maybe_compress,
|
||||
update_kpool_write_plan_cuda_graph,
|
||||
)
|
||||
from sglang.srt.layers.attention.dsa.kpool_plan import (
|
||||
_alloc_kpool_write_plan_buffers,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
|
||||
|
||||
@unittest.skipUnless(torch.cuda.is_available(), "Test requires CUDA")
|
||||
class TestDsaKpoolMultiPool(CustomTestCase):
|
||||
POOL_SIZE = 4
|
||||
PAGE_SIZE = 64
|
||||
SLOTS_PER_PAGE = 64
|
||||
NUM_DRAFT_TOKENS = 6
|
||||
|
||||
def _pool(self) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
page_size=self.PAGE_SIZE,
|
||||
index_head_dim=INDEX_HEAD_DIM,
|
||||
slots_per_page=self.SLOTS_PER_PAGE,
|
||||
index_kpool=self.POOL_SIZE,
|
||||
tail_extra_slots=self.NUM_DRAFT_TOKENS,
|
||||
quant_block_size=128,
|
||||
)
|
||||
|
||||
def _empty_cache(self) -> torch.Tensor:
|
||||
page_nbytes = self.SLOTS_PER_PAGE * INDEX_HEAD_DIM + self.SLOTS_PER_PAGE * 4
|
||||
return torch.zeros((1, page_nbytes), dtype=torch.uint8, device="cuda")
|
||||
|
||||
def test_write_plan_records_every_candidate_pool(self):
|
||||
batch_size = 2
|
||||
num_draft_tokens = self.NUM_DRAFT_TOKENS
|
||||
max_closed_pools = kpool_max_closed_pools(num_draft_tokens, self.POOL_SIZE)
|
||||
self.assertEqual(max_closed_pools, 2)
|
||||
|
||||
plan = _alloc_kpool_write_plan_buffers(
|
||||
max_bs=batch_size,
|
||||
num_draft_tokens=num_draft_tokens,
|
||||
pool_size=self.POOL_SIZE,
|
||||
device=torch.device("cuda"),
|
||||
is_verify=True,
|
||||
)
|
||||
self.assertEqual(plan.write_loc.shape, (batch_size, max_closed_pools))
|
||||
|
||||
write_start = torch.tensor([3, 255], dtype=torch.int32, device="cuda")
|
||||
req_pool_indices = torch.tensor([7, 11], dtype=torch.int64, device="cuda")
|
||||
real_page_table = torch.zeros(
|
||||
(batch_size * num_draft_tokens, 8),
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
real_page_table[:num_draft_tokens, 0] = 2
|
||||
real_page_table[:num_draft_tokens, 4] = 3
|
||||
real_page_table[num_draft_tokens:, 0] = 5
|
||||
real_page_table[num_draft_tokens:, 4] = 6
|
||||
|
||||
update_kpool_write_plan_cuda_graph(
|
||||
write_start=write_start,
|
||||
req_pool_indices=req_pool_indices,
|
||||
real_page_table=real_page_table,
|
||||
req_out=plan.req,
|
||||
write_start_out=plan.write_start,
|
||||
tail_logical_start_out=plan.tail_logical_start,
|
||||
write_loc_out=plan.write_loc,
|
||||
pool_seqlens_per_q_out=plan.pool_seqlens_per_q,
|
||||
seqlens_per_q_out=plan.seqlens_per_q,
|
||||
pool_size=self.POOL_SIZE,
|
||||
num_draft_tokens=num_draft_tokens,
|
||||
slots_per_page=self.SLOTS_PER_PAGE,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(plan.req, req_pool_indices)
|
||||
torch.testing.assert_close(plan.write_start, write_start)
|
||||
torch.testing.assert_close(
|
||||
plan.tail_logical_start,
|
||||
torch.tensor([0, 252], dtype=torch.int32, device="cuda"),
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
plan.write_loc,
|
||||
torch.tensor(
|
||||
[
|
||||
[2 * self.SLOTS_PER_PAGE, 2 * self.SLOTS_PER_PAGE + 1],
|
||||
[
|
||||
5 * self.SLOTS_PER_PAGE + 63,
|
||||
6 * self.SLOTS_PER_PAGE,
|
||||
],
|
||||
],
|
||||
dtype=torch.int64,
|
||||
device="cuda",
|
||||
),
|
||||
)
|
||||
|
||||
def _run_compress_case(self, effective_n: int, expected_closed_pools: int):
|
||||
torch.manual_seed(42)
|
||||
pool = self._pool()
|
||||
num_draft_tokens = self.NUM_DRAFT_TOKENS
|
||||
tail_size = self.POOL_SIZE + num_draft_tokens
|
||||
write_start_value = 3
|
||||
|
||||
key = torch.randn(
|
||||
num_draft_tokens, INDEX_HEAD_DIM, dtype=torch.bfloat16, device="cuda"
|
||||
)
|
||||
score = torch.randn_like(key)
|
||||
ape = torch.randn(
|
||||
self.POOL_SIZE, INDEX_HEAD_DIM, dtype=torch.float32, device="cuda"
|
||||
)
|
||||
tail_k_initial = torch.randn(
|
||||
1, tail_size, INDEX_HEAD_DIM, dtype=torch.bfloat16, device="cuda"
|
||||
)
|
||||
tail_score_initial = torch.randn_like(tail_k_initial)
|
||||
|
||||
tail_k_expected = tail_k_initial.clone()
|
||||
tail_score_expected = tail_score_initial.clone()
|
||||
for i in range(num_draft_tokens):
|
||||
physical_slot = (write_start_value + i) % tail_size
|
||||
tail_k_expected[0, physical_slot] = key[i]
|
||||
tail_score_expected[0, physical_slot] = score[i]
|
||||
|
||||
expected_cache = self._empty_cache()
|
||||
dummy_chunk = torch.zeros(
|
||||
1, INDEX_HEAD_DIM, dtype=torch.bfloat16, device="cuda"
|
||||
)
|
||||
kpool_assemble_softmax_rotate_write_cache(
|
||||
pool=pool,
|
||||
buf=expected_cache,
|
||||
chunk_k=dummy_chunk,
|
||||
chunk_score=dummy_chunk,
|
||||
tail_k=tail_k_expected,
|
||||
tail_score=tail_score_expected,
|
||||
req_pool_idx=torch.zeros(
|
||||
expected_closed_pools, dtype=torch.int64, device="cuda"
|
||||
),
|
||||
n_from_tail=torch.full(
|
||||
(expected_closed_pools,),
|
||||
self.POOL_SIZE,
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
),
|
||||
chunk_src_start=torch.zeros(
|
||||
expected_closed_pools, dtype=torch.int64, device="cuda"
|
||||
),
|
||||
tail_logical_base=torch.arange(
|
||||
0,
|
||||
expected_closed_pools * self.POOL_SIZE,
|
||||
self.POOL_SIZE,
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
),
|
||||
ape=ape,
|
||||
loc=torch.arange(expected_closed_pools, dtype=torch.int64, device="cuda"),
|
||||
round_scale=False,
|
||||
)
|
||||
|
||||
actual_cache = self._empty_cache()
|
||||
tail_k_actual = tail_k_initial.clone()
|
||||
tail_score_actual = tail_score_initial.clone()
|
||||
kpool_write_tail_and_maybe_compress(
|
||||
pool=pool,
|
||||
buf=actual_cache,
|
||||
key=key,
|
||||
score=score,
|
||||
tail_k=tail_k_actual,
|
||||
tail_score=tail_score_actual,
|
||||
ape=ape,
|
||||
req_pool_indices=torch.zeros(1, dtype=torch.int64, device="cuda"),
|
||||
write_start=torch.tensor(
|
||||
[write_start_value], dtype=torch.int32, device="cuda"
|
||||
),
|
||||
tail_logical_start=torch.zeros(1, dtype=torch.int32, device="cuda"),
|
||||
write_loc=torch.tensor([[0, 1]], dtype=torch.int64, device="cuda"),
|
||||
out_cache_loc=torch.arange(
|
||||
1, num_draft_tokens + 1, dtype=torch.int64, device="cuda"
|
||||
),
|
||||
num_draft_tokens=num_draft_tokens,
|
||||
round_scale=False,
|
||||
effective_n_per_batch=torch.tensor(
|
||||
[effective_n], dtype=torch.int32, device="cuda"
|
||||
),
|
||||
)
|
||||
|
||||
torch.testing.assert_close(tail_k_actual, tail_k_expected, atol=0, rtol=0)
|
||||
torch.testing.assert_close(
|
||||
tail_score_actual, tail_score_expected, atol=0, rtol=0
|
||||
)
|
||||
torch.testing.assert_close(actual_cache, expected_cache, atol=0, rtol=0)
|
||||
|
||||
def test_compresses_two_pools_when_draft_window_closes_two(self):
|
||||
self._run_compress_case(
|
||||
effective_n=self.NUM_DRAFT_TOKENS,
|
||||
expected_closed_pools=2,
|
||||
)
|
||||
|
||||
def test_effective_n_only_compresses_accepted_pools(self):
|
||||
self._run_compress_case(effective_n=2, expected_closed_pools=1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user