[Attention][MLA] Fuse Kimi-K3 chunked-context K/V packing (#51772)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Yongye Zhu
co-authored by
Claude Opus 5
mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Claude Opus 5
mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent
63344913a6
commit
903d2efe7e
@@ -25,6 +25,14 @@
|
||||
* matching MLA's _q_scale / _k_scale / _v_scale (the cache latent uses a
|
||||
* separate cache scale). Per-tensor E4M3.
|
||||
*
|
||||
* K/V pack (fused_kimi_k3_mla_kv_concat{,_quant_fp8}):
|
||||
* - k_out[t, h] = [k_nope[t, h] | k_pe[t]], cast to fp8 by the _quant_fp8
|
||||
* variant, which also writes v_fp8[t, h] = cast(v[t, h])
|
||||
* The chunked-context epilogue: no cache insert, no q, no RoPE (the cached
|
||||
* k_pe is already rotated) and no scaling. `k_nope`/`v` are strided views
|
||||
* of one kv_b_proj output and `k_pe` may already be fp8 (a plain fp8 cache is
|
||||
* gathered without dequantizing); the outputs are contiguous.
|
||||
*
|
||||
* fp8_ds_mla (fused_kimi_k3_mla_key_concat_ds_mla_insert):
|
||||
* - full key concat (bf16), and cache insert in DeepSeek's 656-byte
|
||||
* block-scaled layout (NoPE fp8 in 4 tiles of 128 with per-tile dynamic
|
||||
@@ -183,6 +191,36 @@ __device__ __forceinline__ void copyChunk8(void* dst, const scalar_t* src,
|
||||
}
|
||||
}
|
||||
|
||||
// Cast 8 source elements (one uint4 of bf16/fp16) to E4M3 with no scaling,
|
||||
// using the native pairwise converters so the result is bit-identical to
|
||||
// `src.to(torch.float8_e4m3fn)`.
|
||||
template <typename scalar_t>
|
||||
__device__ __forceinline__ void copyChunk8UnitFp8(uint8_t* dst,
|
||||
const scalar_t* src) {
|
||||
#ifndef USE_ROCM
|
||||
uint4 const input = *reinterpret_cast<const uint4*>(src);
|
||||
using Converter = vllm::_typeConvert<scalar_t>;
|
||||
auto const* input2 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type const*>(&input);
|
||||
uint2 output;
|
||||
auto* output2 = reinterpret_cast<__nv_fp8x2_storage_t*>(&output);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
if constexpr (std::is_same_v<scalar_t, c10::BFloat16>) {
|
||||
output2[i] = __nv_cvt_bfloat16raw2_to_fp8x2(
|
||||
static_cast<__nv_bfloat162_raw>(input2[i]), __NV_SATFINITE,
|
||||
__NV_E4M3);
|
||||
} else {
|
||||
output2[i] = __nv_cvt_halfraw2_to_fp8x2(
|
||||
static_cast<__half2_raw>(input2[i]), __NV_SATFINITE, __NV_E4M3);
|
||||
}
|
||||
}
|
||||
*reinterpret_cast<uint2*>(dst) = output;
|
||||
#else
|
||||
copyChunk8<scalar_t, true>(dst, src, 1.0f);
|
||||
#endif
|
||||
}
|
||||
|
||||
// Concat + store one head's full key: dst[e] = [k_nope | k_pe], e in [0, 192).
|
||||
// FP8 dst is byte-addressed; bf16 dst is scalar_t-addressed (dst_elem_size).
|
||||
template <typename scalar_t, bool FP8, bool APPLY_ROPE = false>
|
||||
@@ -202,6 +240,38 @@ __device__ __forceinline__ void writeFullKey(void* dst, const scalar_t* k_nope,
|
||||
}
|
||||
}
|
||||
|
||||
// Unscaled full-key writer for the chunked-context pack, driven by a half warp
|
||||
// (16 lanes x 8 elems covers NoPE 128; 8 of them cover RoPE 64). Unlike
|
||||
// `writeFullKey` there is no scale and no RoPE (a gathered k_pe is already
|
||||
// rotated), and KPE_FP8 handles a k_pe that the gather left in the fp8 cache
|
||||
// layout -- so neither dtype needs an intermediate cast or a broadcast copy.
|
||||
template <typename scalar_t, bool FP8, bool KPE_FP8>
|
||||
__device__ __forceinline__ void writeFullKeyPack(void* dst,
|
||||
const scalar_t* k_nope,
|
||||
const void* k_pe, int laneId) {
|
||||
static_assert(!KPE_FP8 || FP8, "an fp8 k_pe requires an fp8 key output");
|
||||
constexpr int kElemSize = FP8 ? 1 : sizeof(scalar_t);
|
||||
auto* d = reinterpret_cast<uint8_t*>(dst);
|
||||
for (int e = laneId * kVecElems; e < kQkHeadDim; e += 16 * kVecElems) {
|
||||
void* out = d + e * kElemSize;
|
||||
const scalar_t* src = k_nope + e;
|
||||
if (e >= kQkNopeHeadDim) {
|
||||
int const rope_e = e - kQkNopeHeadDim;
|
||||
if constexpr (KPE_FP8) {
|
||||
*reinterpret_cast<uint2*>(out) = *reinterpret_cast<const uint2*>(
|
||||
reinterpret_cast<const uint8_t*>(k_pe) + rope_e);
|
||||
continue;
|
||||
}
|
||||
src = reinterpret_cast<const scalar_t*>(k_pe) + rope_e;
|
||||
}
|
||||
if constexpr (FP8) {
|
||||
copyChunk8UnitFp8(reinterpret_cast<uint8_t*>(out), src);
|
||||
} else {
|
||||
*reinterpret_cast<uint4*>(out) = *reinterpret_cast<const uint4*>(src);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Store a prefill query, rotating only q[..., 128:192]. For bf16 dst may alias
|
||||
// q (in-place); fp8 writes the quantized query directly to its output.
|
||||
template <typename scalar_t, bool FP8, bool APPLY_ROPE = false>
|
||||
@@ -457,6 +527,72 @@ __global__ void fusedKimiK3MLAQKVQuantKVCacheFp8Kernel(
|
||||
#endif
|
||||
}
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// K/V pack variant (chunked context): concatenate the strided kv_b_proj output
|
||||
// with k_pe into a contiguous row-major key, casting to E4M3 (and casting V
|
||||
// into its own contiguous output) when FP8. Unlike the other kernels this one
|
||||
// has no per-token cache slot, so a warp handles two (token, head) rows -- 16
|
||||
// lanes each -- and the grid is capped so a long context becomes a grid-stride
|
||||
// loop instead of a huge launch.
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
template <typename scalar_t, bool FP8, bool KPE_FP8>
|
||||
__global__ void fusedKimiK3MLAKVConcatPackKernel(
|
||||
const scalar_t* __restrict__ k_nope, int64_t const kn_tok_stride,
|
||||
int64_t const kn_head_stride, const void* __restrict__ k_pe,
|
||||
int64_t const k_pe_tok_stride, const scalar_t* __restrict__ v,
|
||||
int64_t const v_tok_stride, int64_t const v_head_stride,
|
||||
void* __restrict__ k_out, int64_t const ko_tok_stride,
|
||||
int64_t const ko_head_stride, uint8_t* __restrict__ v_fp8,
|
||||
int64_t const vo_tok_stride, int64_t const vo_head_stride,
|
||||
int const num_tokens, int const num_heads) {
|
||||
int const warpsPerBlock = blockDim.x / 32;
|
||||
int const laneId = threadIdx.x % 16;
|
||||
int const rowInWarp = (threadIdx.x % 32) / 16;
|
||||
int64_t const globalWarpIdx =
|
||||
static_cast<int64_t>(blockIdx.x) * warpsPerBlock + threadIdx.x / 32;
|
||||
int64_t const totalRows = static_cast<int64_t>(num_tokens) * num_heads;
|
||||
if (globalWarpIdx * 2 >= totalRows) return;
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaGridDependencySynchronize();
|
||||
#endif
|
||||
|
||||
constexpr int kOutElemSize = FP8 ? 1 : sizeof(scalar_t);
|
||||
int64_t const gridStride =
|
||||
static_cast<int64_t>(gridDim.x) * warpsPerBlock * 2;
|
||||
for (int64_t row = globalWarpIdx * 2 + rowInWarp; row < totalRows;
|
||||
row += gridStride) {
|
||||
int const tokenIdx = static_cast<int>(row / num_heads);
|
||||
int const headIdx = static_cast<int>(row % num_heads);
|
||||
const void* pe;
|
||||
if constexpr (KPE_FP8) {
|
||||
pe = reinterpret_cast<const uint8_t*>(k_pe) + tokenIdx * k_pe_tok_stride;
|
||||
} else {
|
||||
pe = reinterpret_cast<const scalar_t*>(k_pe) + tokenIdx * k_pe_tok_stride;
|
||||
}
|
||||
writeFullKeyPack<scalar_t, FP8, KPE_FP8>(
|
||||
reinterpret_cast<uint8_t*>(k_out) +
|
||||
(tokenIdx * ko_tok_stride + headIdx * ko_head_stride) *
|
||||
kOutElemSize,
|
||||
k_nope + tokenIdx * kn_tok_stride + headIdx * kn_head_stride, pe,
|
||||
laneId);
|
||||
|
||||
// bf16 V needs no cast, so it stays a strided view of the kv_b_proj output.
|
||||
if constexpr (FP8) {
|
||||
const scalar_t* vh =
|
||||
v + tokenIdx * v_tok_stride + headIdx * v_head_stride;
|
||||
uint8_t* vo = v_fp8 + tokenIdx * vo_tok_stride + headIdx * vo_head_stride;
|
||||
for (int e = laneId * kVecElems; e < kVHeadDim; e += 16 * kVecElems) {
|
||||
copyChunk8UnitFp8(vo + e, vh + e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
#endif
|
||||
}
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// ds_mla variant: concat full key (bf16) + fp8_ds_mla latent cache insert
|
||||
//
|
||||
@@ -648,16 +784,20 @@ __global__ void fusedKimiK3MLADecodeQConcatDsMlaKernel(
|
||||
#endif
|
||||
}
|
||||
|
||||
// PDL-aware launch of a (token, num_heads + 1)-warp grid.
|
||||
// PDL-aware launch of a (token, slots_per_token)-row grid, `rows_per_warp` rows
|
||||
// to a warp. `max_blocks > 0` caps the grid, leaving the kernel's grid-stride
|
||||
// loop to cover the remaining rows.
|
||||
template <typename KernelT, typename... Args>
|
||||
static void launchPdl(KernelT kernel, int num_tokens, int num_heads,
|
||||
cudaStream_t stream, Args... args) {
|
||||
static void launchPdlSlots(KernelT kernel, int num_tokens, int slots_per_token,
|
||||
int rows_per_warp, int max_blocks,
|
||||
cudaStream_t stream, Args... args) {
|
||||
constexpr int kBlockSize = 256;
|
||||
constexpr int kWarpsPerBlock = kBlockSize / 32;
|
||||
int64_t const total_warps =
|
||||
static_cast<int64_t>(num_tokens) * (num_heads + 1);
|
||||
int const grid =
|
||||
int64_t const total_rows = static_cast<int64_t>(num_tokens) * slots_per_token;
|
||||
int64_t const total_warps = (total_rows + rows_per_warp - 1) / rows_per_warp;
|
||||
int grid =
|
||||
static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
|
||||
if (max_blocks > 0 && grid > max_blocks) grid = max_blocks;
|
||||
#ifndef USE_ROCM
|
||||
static int const sm_version = getSMVersion();
|
||||
cudaLaunchConfig_t config;
|
||||
@@ -681,6 +821,13 @@ static void launchPdl(KernelT kernel, int num_tokens, int num_heads,
|
||||
#endif
|
||||
}
|
||||
|
||||
// Cache-inserting kernels use one extra warp per token for the cache slot.
|
||||
template <typename KernelT, typename... Args>
|
||||
static void launchPdl(KernelT kernel, int num_tokens, int num_heads,
|
||||
cudaStream_t stream, Args... args) {
|
||||
launchPdlSlots(kernel, num_tokens, num_heads + 1, 1, 0, stream, args...);
|
||||
}
|
||||
|
||||
void checkBfloat16Support(torch::headeronly::ScalarType dtype) {
|
||||
#ifndef USE_ROCM
|
||||
if (dtype == torch::headeronly::ScalarType::BFloat16) {
|
||||
@@ -913,6 +1060,128 @@ void fused_kimi_k3_mla_key_concat_ds_mla_insert(
|
||||
});
|
||||
}
|
||||
|
||||
namespace {
|
||||
// Shared checks for the chunked-context K/V pack ops. `k_nope` / `v` are
|
||||
// strided views of one kv_b_proj output and `k_pe` is a strided view of the
|
||||
// gathered context workspace, so only a unit last-dim stride (not full
|
||||
// contiguity) is required of the inputs. Returns true when `k_pe` is still in
|
||||
// the fp8 cache layout.
|
||||
bool check_kv_concat_inputs(torch::stable::Tensor const& k_nope,
|
||||
torch::stable::Tensor const& k_pe,
|
||||
torch::stable::Tensor const& k_out,
|
||||
torch::headeronly::ScalarType out_dtype) {
|
||||
using torch::headeronly::ScalarType;
|
||||
ScalarType const dt = k_nope.scalar_type();
|
||||
STD_TORCH_CHECK(k_nope.device().is_cuda() && k_nope.dim() == 3 &&
|
||||
k_nope.size(2) == 128 && k_nope.stride(2) == 1,
|
||||
"k_nope must be CUDA [T, H, 128] with unit last stride");
|
||||
bool const k_pe_is_fp8 = k_pe.scalar_type() == ScalarType::Float8_e4m3fn;
|
||||
STD_TORCH_CHECK(
|
||||
k_pe.device() == k_nope.device() && k_pe.dim() == 2 &&
|
||||
k_pe.size(1) == 64 && k_pe.stride(1) == 1 &&
|
||||
(k_pe.scalar_type() == dt ||
|
||||
(k_pe_is_fp8 && out_dtype == ScalarType::Float8_e4m3fn)),
|
||||
"k_pe must be CUDA [T, 64] with unit last-dim stride and use the input "
|
||||
"dtype, or float8_e4m3fn when the key output is fp8");
|
||||
STD_TORCH_CHECK(k_out.device() == k_nope.device() && k_out.is_contiguous() &&
|
||||
k_out.scalar_type() == out_dtype && k_out.dim() == 3 &&
|
||||
k_out.size(2) == 192 && k_out.size(0) == k_nope.size(0) &&
|
||||
k_out.size(1) == k_nope.size(1) &&
|
||||
k_pe.size(0) == k_nope.size(0),
|
||||
"k_out must be a contiguous CUDA [T, H, 192] tensor of the "
|
||||
"output dtype matching k_nope's token and head dimensions");
|
||||
vllm::kimi_k3_fused_ops::checkBfloat16Support(dt);
|
||||
return k_pe_is_fp8;
|
||||
}
|
||||
} // namespace
|
||||
|
||||
void fused_kimi_k3_mla_kv_concat(
|
||||
torch::stable::Tensor const& k_nope, // [T, H, 128] bf16/fp16
|
||||
torch::stable::Tensor const& k_pe, // [T, 64] same dtype as k_nope
|
||||
torch::stable::Tensor& k_out) { // [T, H, 192] same dtype, written
|
||||
namespace kk3 = vllm::kimi_k3_fused_ops;
|
||||
torch::headeronly::ScalarType const dt = k_nope.scalar_type();
|
||||
check_kv_concat_inputs(k_nope, k_pe, k_out, dt);
|
||||
|
||||
int const num_tokens = static_cast<int>(k_nope.size(0));
|
||||
int const num_heads = static_cast<int>(k_nope.size(1));
|
||||
if (num_tokens == 0) return;
|
||||
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
k_nope.get_device_index());
|
||||
const cudaStream_t stream =
|
||||
get_current_cuda_stream(k_nope.get_device_index());
|
||||
|
||||
VLLM_STABLE_DISPATCH_HALF_TYPES(dt, "fused_kimi_k3_mla_kv_concat", [&] {
|
||||
kk3::launchPdlSlots(
|
||||
kk3::fusedKimiK3MLAKVConcatPackKernel<scalar_t, false, false>,
|
||||
num_tokens, num_heads, 2, get_device_prop()->multiProcessorCount * 8,
|
||||
stream, reinterpret_cast<scalar_t const*>(k_nope.const_data_ptr()),
|
||||
k_nope.stride(0), k_nope.stride(1), k_pe.const_data_ptr(),
|
||||
k_pe.stride(0), static_cast<scalar_t const*>(nullptr), int64_t{0},
|
||||
int64_t{0}, k_out.mutable_data_ptr(), k_out.stride(0), k_out.stride(1),
|
||||
static_cast<uint8_t*>(nullptr), int64_t{0}, int64_t{0}, num_tokens,
|
||||
num_heads);
|
||||
});
|
||||
}
|
||||
|
||||
void fused_kimi_k3_mla_kv_concat_quant_fp8(
|
||||
torch::stable::Tensor const& k_nope, // [T, H, 128] bf16/fp16
|
||||
torch::stable::Tensor const& k_pe, // [T, 64] bf16/fp16/fp8
|
||||
torch::stable::Tensor const& v, // [T, H, 128] bf16/fp16
|
||||
torch::stable::Tensor& k_fp8, // [T, H, 192] fp8, written
|
||||
torch::stable::Tensor& v_fp8) { // [T, H, 128] fp8, written
|
||||
using torch::headeronly::ScalarType;
|
||||
namespace kk3 = vllm::kimi_k3_fused_ops;
|
||||
ScalarType const dt = k_nope.scalar_type();
|
||||
bool const k_pe_is_fp8 =
|
||||
check_kv_concat_inputs(k_nope, k_pe, k_fp8, ScalarType::Float8_e4m3fn);
|
||||
STD_TORCH_CHECK(v.device() == k_nope.device() && v.scalar_type() == dt &&
|
||||
v.dim() == 3 && v.size(2) == 128 && v.stride(2) == 1 &&
|
||||
v.size(0) == k_nope.size(0) &&
|
||||
v.size(1) == k_nope.size(1),
|
||||
"v must match k_nope's dtype and token/head dimensions and "
|
||||
"have shape [T, H, 128] with unit last stride");
|
||||
STD_TORCH_CHECK(v_fp8.device() == k_nope.device() && v_fp8.is_contiguous() &&
|
||||
v_fp8.scalar_type() == ScalarType::Float8_e4m3fn &&
|
||||
v_fp8.dim() == 3 && v_fp8.size(2) == 128 &&
|
||||
v_fp8.size(0) == k_nope.size(0) &&
|
||||
v_fp8.size(1) == k_nope.size(1),
|
||||
"v_fp8 must be a contiguous CUDA [T, H, 128] fp8 tensor "
|
||||
"matching k_nope's token and head dimensions");
|
||||
|
||||
int const num_tokens = static_cast<int>(k_nope.size(0));
|
||||
int const num_heads = static_cast<int>(k_nope.size(1));
|
||||
if (num_tokens == 0) return;
|
||||
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
k_nope.get_device_index());
|
||||
const cudaStream_t stream =
|
||||
get_current_cuda_stream(k_nope.get_device_index());
|
||||
|
||||
VLLM_STABLE_DISPATCH_HALF_TYPES(
|
||||
dt, "fused_kimi_k3_mla_kv_concat_quant_fp8", [&] {
|
||||
auto launch = [&](auto kernel) {
|
||||
kk3::launchPdlSlots(
|
||||
kernel, num_tokens, num_heads, 2,
|
||||
get_device_prop()->multiProcessorCount * 8, stream,
|
||||
reinterpret_cast<scalar_t const*>(k_nope.const_data_ptr()),
|
||||
k_nope.stride(0), k_nope.stride(1), k_pe.const_data_ptr(),
|
||||
k_pe.stride(0),
|
||||
reinterpret_cast<scalar_t const*>(v.const_data_ptr()),
|
||||
v.stride(0), v.stride(1), k_fp8.mutable_data_ptr(),
|
||||
k_fp8.stride(0), k_fp8.stride(1),
|
||||
reinterpret_cast<uint8_t*>(v_fp8.mutable_data_ptr()),
|
||||
v_fp8.stride(0), v_fp8.stride(1), num_tokens, num_heads);
|
||||
};
|
||||
if (k_pe_is_fp8) {
|
||||
launch(kk3::fusedKimiK3MLAKVConcatPackKernel<scalar_t, true, true>);
|
||||
} else {
|
||||
launch(kk3::fusedKimiK3MLAKVConcatPackKernel<scalar_t, true, false>);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(
|
||||
torch::stable::Tensor const& q, // [Tp, H, 192] bf16
|
||||
torch::stable::Tensor const& k_nope, // [Tp, H, 128] bf16
|
||||
|
||||
@@ -300,6 +300,16 @@ void fused_kimi_k3_mla_key_concat_ds_mla_insert(
|
||||
std::optional<torch::stable::Tensor> position_ids,
|
||||
std::optional<torch::stable::Tensor> cos_sin_cache);
|
||||
|
||||
void fused_kimi_k3_mla_kv_concat(torch::stable::Tensor const& k_nope,
|
||||
torch::stable::Tensor const& k_pe,
|
||||
torch::stable::Tensor& k_out);
|
||||
|
||||
void fused_kimi_k3_mla_kv_concat_quant_fp8(torch::stable::Tensor const& k_nope,
|
||||
torch::stable::Tensor const& k_pe,
|
||||
torch::stable::Tensor const& v,
|
||||
torch::stable::Tensor& k_fp8,
|
||||
torch::stable::Tensor& v_fp8);
|
||||
|
||||
void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(
|
||||
torch::stable::Tensor const& q, torch::stable::Tensor const& k_nope,
|
||||
torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed,
|
||||
|
||||
@@ -466,6 +466,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
"Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, "
|
||||
"int cache_block_size, Tensor? position_ids=None, "
|
||||
"Tensor? cos_sin_cache=None) -> ()");
|
||||
ops.def(
|
||||
"fused_kimi_k3_mla_kv_concat(Tensor k_nope, Tensor k_pe, Tensor! k_out) "
|
||||
"-> ()");
|
||||
ops.def(
|
||||
"fused_kimi_k3_mla_kv_concat_quant_fp8("
|
||||
"Tensor k_nope, Tensor k_pe, Tensor v, Tensor! k_fp8, Tensor! v_fp8) "
|
||||
"-> ()");
|
||||
ops.def(
|
||||
"fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert("
|
||||
"Tensor q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, Tensor v, "
|
||||
@@ -769,6 +776,10 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
|
||||
TORCH_BOX(&fused_kimi_k3_mla_key_concat_kv_cache_insert));
|
||||
ops.impl("fused_kimi_k3_mla_key_concat_ds_mla_insert",
|
||||
TORCH_BOX(&fused_kimi_k3_mla_key_concat_ds_mla_insert));
|
||||
ops.impl("fused_kimi_k3_mla_kv_concat",
|
||||
TORCH_BOX(&fused_kimi_k3_mla_kv_concat));
|
||||
ops.impl("fused_kimi_k3_mla_kv_concat_quant_fp8",
|
||||
TORCH_BOX(&fused_kimi_k3_mla_kv_concat_quant_fp8));
|
||||
ops.impl("fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert",
|
||||
TORCH_BOX(&fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert));
|
||||
ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_insert",
|
||||
|
||||
@@ -9,6 +9,8 @@ from vllm.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import (
|
||||
fused_mla_decode_q_concat_kv_cache_insert,
|
||||
fused_mla_key_concat_ds_mla_insert,
|
||||
fused_mla_key_concat_kv_cache_insert,
|
||||
fused_mla_kv_concat,
|
||||
fused_mla_kv_concat_quant_fp8,
|
||||
fused_mla_qkv_quant_kv_cache_fp8_insert,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
@@ -25,8 +27,8 @@ _POSITIONS = (1, 7, 13)
|
||||
_SLOTS = (0, 3, 9)
|
||||
|
||||
|
||||
def _randn(*shape: int) -> torch.Tensor:
|
||||
return torch.randn(*shape, device="cuda", dtype=_DTYPE) * 0.2
|
||||
def _randn(*shape: int, dtype: torch.dtype = _DTYPE) -> torch.Tensor:
|
||||
return torch.randn(*shape, device="cuda", dtype=dtype) * 0.2
|
||||
|
||||
|
||||
def _rope_cache(max_position: int = 32) -> torch.Tensor:
|
||||
@@ -66,6 +68,70 @@ def _assert_fp8_close(actual: torch.Tensor, expected: torch.Tensor) -> None:
|
||||
)
|
||||
|
||||
|
||||
def _strided_context_inputs(
|
||||
num_tokens: int, num_heads: int, dtype: torch.dtype, k_pe_fp8: bool
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""The real prefill-context layouts: ``k_nope``/``v`` as the two strided
|
||||
halves of one kv_b_proj output and ``k_pe`` as a column slice of the gather
|
||||
workspace (left in the fp8 cache layout for a plain fp8 cache)."""
|
||||
kv_nope = _randn(num_tokens, num_heads, 256, dtype=dtype)
|
||||
k_nope, v = kv_nope.split((128, 128), dim=-1)
|
||||
workspace = _randn(num_tokens, 576, dtype=dtype)
|
||||
if k_pe_fp8:
|
||||
workspace = workspace.to(torch.float8_e4m3fn)
|
||||
return k_nope, workspace[:, 512:].unsqueeze(1), v
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [0, _NUM_TOKENS])
|
||||
@pytest.mark.parametrize("num_heads", [3, _NUM_HEADS])
|
||||
@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float16])
|
||||
@torch.inference_mode()
|
||||
def test_context_kv_concat_accepts_strided_inputs(
|
||||
num_tokens: int, num_heads: int, input_dtype: torch.dtype
|
||||
) -> None:
|
||||
k_nope, k_pe, _ = _strided_context_inputs(
|
||||
num_tokens, num_heads, input_dtype, k_pe_fp8=False
|
||||
)
|
||||
|
||||
k = fused_mla_kv_concat(k_nope, k_pe)
|
||||
|
||||
assert k.is_contiguous()
|
||||
assert k.dtype == input_dtype
|
||||
torch.testing.assert_close(k[..., :128], k_nope, atol=0, rtol=0)
|
||||
torch.testing.assert_close(
|
||||
k[..., 128:], k_pe.expand(-1, num_heads, -1), atol=0, rtol=0
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("k_pe_dtype", ["input", "fp8"])
|
||||
@pytest.mark.parametrize("num_tokens", [0, _NUM_TOKENS])
|
||||
@pytest.mark.parametrize("num_heads", [3, _NUM_HEADS])
|
||||
@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float16])
|
||||
@torch.inference_mode()
|
||||
def test_context_kv_pack_quantizes_strided_inputs(
|
||||
k_pe_dtype: str,
|
||||
num_tokens: int,
|
||||
num_heads: int,
|
||||
input_dtype: torch.dtype,
|
||||
) -> None:
|
||||
k_nope, k_pe, v = _strided_context_inputs(
|
||||
num_tokens, num_heads, input_dtype, k_pe_fp8=k_pe_dtype == "fp8"
|
||||
)
|
||||
|
||||
k_actual, v_actual = fused_mla_kv_concat_quant_fp8(k_nope, k_pe, v)
|
||||
|
||||
fp8 = torch.float8_e4m3fn
|
||||
k_expected = torch.empty_like(k_actual)
|
||||
k_expected[..., :128] = k_nope.to(fp8)
|
||||
k_expected[..., 128:] = k_pe.to(fp8)
|
||||
assert k_actual.is_contiguous()
|
||||
assert v_actual.is_contiguous()
|
||||
# The pack casts with the native pairwise converters, so it must be
|
||||
# bit-identical to torch's `.to(fp8)` rather than merely close.
|
||||
torch.testing.assert_close(k_actual.float(), k_expected.float(), atol=0, rtol=0)
|
||||
torch.testing.assert_close(v_actual.float(), v.to(fp8).float(), atol=0, rtol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cache_kind", ["bf16", "fp8", "fp8_ds_mla"])
|
||||
@torch.inference_mode()
|
||||
def test_prefill_epilogue_fuses_gptj_rope(cache_kind: str) -> None:
|
||||
|
||||
@@ -0,0 +1,319 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""K3's fused chunked-context prefill must match the generic MLA impl.
|
||||
|
||||
The layer owns its own context loop so it can fuse the per-chunk K/V pack and
|
||||
skip re-quantizing an already-quantized query. That is only safe if it feeds the
|
||||
prefill backend exactly what ``MLACommonBaseImpl._compute_prefill_context``
|
||||
would, chunk for chunk.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.attention.mla_attention import (
|
||||
MLACommonBaseImpl,
|
||||
MLACommonPrefillMetadata,
|
||||
build_mla_chunked_context_metadata,
|
||||
)
|
||||
from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not current_platform.is_cuda(), reason="Kimi-K3 fused MLA requires CUDA"
|
||||
)
|
||||
|
||||
_KV_LORA_RANK = 512
|
||||
_QK_NOPE = 128
|
||||
_QK_ROPE = 64
|
||||
_V_HEAD_DIM = 128
|
||||
_ENTRY = _KV_LORA_RANK + _QK_ROPE
|
||||
_NUM_HEADS = 2
|
||||
_BLOCK_SIZE = 16
|
||||
_WORKSPACE_TOKENS = 128
|
||||
# Splits one long request across chunks, packs short ones together, and leaves
|
||||
# the last request without any context.
|
||||
_CONTEXT_LENS = [200, 48, 32, 0]
|
||||
_QUERY_LENS = [8, 4, 6, 5]
|
||||
|
||||
|
||||
class _RecordingPrefillBackend:
|
||||
"""Records what each chunk is asked to attend over.
|
||||
|
||||
``honors_out`` mimics the two backend families: one that writes into a
|
||||
caller-provided ``out`` (trtllm_ragged, flashinfer, tokenspeed) and one that
|
||||
always returns its own buffer (flash_attn with a padded V, aiter).
|
||||
"""
|
||||
|
||||
def __init__(self, honors_out: bool = False) -> None:
|
||||
self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
|
||||
self.out_destinations: list[torch.Tensor | None] = []
|
||||
self._honors_out = honors_out
|
||||
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "recording"
|
||||
|
||||
def supports_out(self) -> bool:
|
||||
return self._honors_out
|
||||
|
||||
def run_prefill_context_chunk(self, *, chunk, q, k, v, out=None):
|
||||
self.calls.append((q.float().clone(), k.float().clone(), v.float().clone()))
|
||||
self.out_destinations.append(out)
|
||||
assert out is None or self._honors_out
|
||||
# Fold K/V into the partial so any packing difference shows up in the
|
||||
# merged context output, not just in the recorded calls.
|
||||
digest = (k.float().mean() + v.float().mean()).item()
|
||||
num_q = q.shape[0]
|
||||
if out is None:
|
||||
out = torch.empty(
|
||||
(num_q, _NUM_HEADS, _V_HEAD_DIM), device=q.device, dtype=torch.bfloat16
|
||||
)
|
||||
else:
|
||||
assert out.shape == (num_q, _NUM_HEADS, _V_HEAD_DIM)
|
||||
out.fill_(digest)
|
||||
lse = torch.full(
|
||||
(_NUM_HEADS, num_q),
|
||||
1.0 + chunk.index,
|
||||
device=q.device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
return out, lse
|
||||
|
||||
|
||||
class _KVBProj(torch.nn.Module):
|
||||
"""Stand-in for the layer's ``kv_b_proj`` (returns an (out, bias) tuple).
|
||||
|
||||
Enforces the same input contract as the real linear methods, because that is
|
||||
what decides whether the gathered latent needs a cast:
|
||||
|
||||
* an fp8 weight consumes the fp8 latent directly and dequantizes internally,
|
||||
so it takes fp8 or bf16;
|
||||
* a bf16 weight -- what a stock K3 checkpoint carries -- is a plain
|
||||
``F.linear`` and rejects anything but bf16, exactly as torch does.
|
||||
|
||||
Either weight dtype produces a bf16 output.
|
||||
"""
|
||||
|
||||
def __init__(self, device: torch.device, weight_dtype: torch.dtype) -> None:
|
||||
super().__init__()
|
||||
weight = (
|
||||
torch.randn(
|
||||
_NUM_HEADS * (_QK_NOPE + _V_HEAD_DIM),
|
||||
_KV_LORA_RANK,
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
* 0.05
|
||||
)
|
||||
self.register_buffer("weight", weight.to(weight_dtype))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, None]:
|
||||
if self.weight.dtype == torch.bfloat16 and x.dtype != torch.bfloat16:
|
||||
raise RuntimeError(
|
||||
"a bfloat16 kv_b_proj cannot consume the gathered latent as "
|
||||
f"{x.dtype}; it must be cast first"
|
||||
)
|
||||
return torch.nn.functional.linear(
|
||||
x.to(torch.bfloat16), self.weight.to(torch.bfloat16)
|
||||
), None
|
||||
|
||||
|
||||
class _FusedLayer:
|
||||
"""Only the attributes K3's context loop reads."""
|
||||
|
||||
_compute_prefill_context = MultiHeadLatentAttention._compute_prefill_context
|
||||
_gather_context_latent = MultiHeadLatentAttention._gather_context_latent
|
||||
_attn_read_kv_cache = MultiHeadLatentAttention._attn_read_kv_cache
|
||||
|
||||
def __init__(self, kv_b_proj, kv_cache, kv_cache_dtype, k_scale) -> None:
|
||||
self.kv_b_proj = kv_b_proj
|
||||
self.kv_cache = kv_cache
|
||||
self.kv_cache_dtype = kv_cache_dtype
|
||||
self._k_scale = k_scale
|
||||
self.kv_lora_rank = _KV_LORA_RANK
|
||||
self.num_local_heads = _NUM_HEADS
|
||||
self.qk_nope_head_dim = _QK_NOPE
|
||||
self.v_head_dim = _V_HEAD_DIM
|
||||
|
||||
|
||||
class _ReferenceImpl:
|
||||
"""Only the attributes the generic context loop reads."""
|
||||
|
||||
_compute_prefill_context = MLACommonBaseImpl._compute_prefill_context
|
||||
_concat_k_nope_k_pe = MLACommonBaseImpl._concat_k_nope_k_pe
|
||||
_use_flashinfer_concat_mla_k = False
|
||||
|
||||
def __init__(self, kv_b_proj, kv_cache_dtype) -> None:
|
||||
self.kv_b_proj = kv_b_proj
|
||||
self.kv_cache_dtype = kv_cache_dtype
|
||||
self.kv_lora_rank = _KV_LORA_RANK
|
||||
self.num_heads = _NUM_HEADS
|
||||
self.qk_nope_head_dim = _QK_NOPE
|
||||
self.qk_rope_head_dim = _QK_ROPE
|
||||
self.v_head_dim = _V_HEAD_DIM
|
||||
|
||||
|
||||
def _build_prefill_metadata(
|
||||
device: torch.device,
|
||||
workspace_dtype: torch.dtype,
|
||||
q_data_type: torch.dtype,
|
||||
backend: _RecordingPrefillBackend,
|
||||
) -> MLACommonPrefillMetadata:
|
||||
query_start_loc_cpu = torch.zeros(len(_QUERY_LENS) + 1, dtype=torch.int32)
|
||||
query_start_loc_cpu[1:] = torch.tensor(_QUERY_LENS, dtype=torch.int32).cumsum(0)
|
||||
workspace = torch.empty(
|
||||
(_WORKSPACE_TOKENS, _ENTRY), dtype=workspace_dtype, device=device
|
||||
)
|
||||
chunked_context = build_mla_chunked_context_metadata(
|
||||
context_lens_cpu=torch.tensor(_CONTEXT_LENS, dtype=torch.int32),
|
||||
prefill_query_start_loc_cpu=query_start_loc_cpu,
|
||||
chunked_prefill_workspace=workspace,
|
||||
chunked_prefill_workspace_size=_WORKSPACE_TOKENS,
|
||||
block_size=_BLOCK_SIZE,
|
||||
align_chunk_to_block=True,
|
||||
device=device,
|
||||
dcp_world_size=1,
|
||||
dcp_local_block_size=1,
|
||||
dcp_virtual_block_size=1,
|
||||
)
|
||||
assert chunked_context is not None
|
||||
assert len(chunked_context.chunks) > 1, "the batch must exercise accumulation"
|
||||
|
||||
max_blocks = (max(_CONTEXT_LENS) + max(_QUERY_LENS)) // _BLOCK_SIZE + 1
|
||||
num_blocks = max_blocks * len(_CONTEXT_LENS)
|
||||
block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).view(
|
||||
len(_CONTEXT_LENS), max_blocks
|
||||
)
|
||||
return MLACommonPrefillMetadata(
|
||||
block_table=block_table,
|
||||
query_start_loc=query_start_loc_cpu.to(device),
|
||||
max_query_len=max(_QUERY_LENS),
|
||||
chunked_context=chunked_context,
|
||||
q_data_type=q_data_type,
|
||||
output_dtype=torch.bfloat16,
|
||||
prefill_backend=backend,
|
||||
), num_blocks
|
||||
|
||||
|
||||
@pytest.mark.parametrize("honors_out", [False, True], ids=["copy_out", "writes_out"])
|
||||
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
|
||||
@pytest.mark.parametrize(
|
||||
"kv_b_proj_quantized", [True, False], ids=["fp8_kv_b_proj", "bf16_kv_b_proj"]
|
||||
)
|
||||
@torch.inference_mode()
|
||||
def test_fused_context_matches_generic_impl(
|
||||
kv_b_proj_quantized: bool, kv_cache_dtype: str, honors_out: bool
|
||||
) -> None:
|
||||
"""Cache dtype and ``kv_b_proj`` dtype vary independently.
|
||||
|
||||
A stock K3 checkpoint pairs a bf16 ``kv_b_proj`` with an fp8 cache, so the
|
||||
fused loop cannot assume the gathered latent is already in the dtype
|
||||
``kv_b_proj`` accepts.
|
||||
"""
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda")
|
||||
fp8 = current_platform.fp8_dtype()
|
||||
quantized = kv_cache_dtype == "fp8"
|
||||
# An fp8 cache is read by an fp8 query, so the workspace keeps the fp8
|
||||
# layout; a bf16 cache dequantizes into a bf16 workspace.
|
||||
q_data_type = fp8 if quantized else torch.bfloat16
|
||||
workspace_dtype = q_data_type
|
||||
|
||||
kv_b_proj = _KVBProj(
|
||||
device, weight_dtype=fp8 if kv_b_proj_quantized else torch.bfloat16
|
||||
)
|
||||
k_scale = torch.ones(1, dtype=torch.float32, device=device)
|
||||
|
||||
backend_fused = _RecordingPrefillBackend(honors_out=honors_out)
|
||||
# The reference impl never passes `out`, so it always allocates its own.
|
||||
backend_ref = _RecordingPrefillBackend()
|
||||
prefill_fused, num_blocks = _build_prefill_metadata(
|
||||
device, workspace_dtype, q_data_type, backend_fused
|
||||
)
|
||||
prefill_ref, _ = _build_prefill_metadata(
|
||||
device, workspace_dtype, q_data_type, backend_ref
|
||||
)
|
||||
|
||||
cache = torch.randn(
|
||||
(num_blocks, _BLOCK_SIZE, _ENTRY), device=device, dtype=torch.bfloat16
|
||||
)
|
||||
kv_cache = cache.to(fp8) if quantized else cache
|
||||
q = (
|
||||
torch.randn(
|
||||
(sum(_QUERY_LENS), _NUM_HEADS, _QK_NOPE + _QK_ROPE),
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
* 0.2
|
||||
).to(q_data_type)
|
||||
|
||||
layer = _FusedLayer(kv_b_proj, kv_cache, kv_cache_dtype, k_scale)
|
||||
fused_out, fused_lse = layer._compute_prefill_context(
|
||||
q, SimpleNamespace(prefill=prefill_fused)
|
||||
)
|
||||
|
||||
impl = _ReferenceImpl(kv_b_proj, kv_cache_dtype)
|
||||
ref_out, ref_lse = impl._compute_prefill_context(
|
||||
q, kv_cache, SimpleNamespace(prefill=prefill_ref), k_scale
|
||||
)
|
||||
|
||||
assert len(backend_fused.calls) == len(backend_ref.calls)
|
||||
for chunk_idx, (fused_call, ref_call) in enumerate(
|
||||
zip(backend_fused.calls, backend_ref.calls, strict=True)
|
||||
):
|
||||
for name, fused_t, ref_t in zip(
|
||||
("q", "k", "v"), fused_call, ref_call, strict=True
|
||||
):
|
||||
torch.testing.assert_close(
|
||||
fused_t,
|
||||
ref_t,
|
||||
atol=0,
|
||||
rtol=0,
|
||||
msg=lambda m, n=name, i=chunk_idx: f"chunk {i} {n} differs: {m}",
|
||||
)
|
||||
torch.testing.assert_close(fused_out, ref_out, atol=0, rtol=0)
|
||||
torch.testing.assert_close(fused_lse, ref_lse, atol=0, rtol=0)
|
||||
|
||||
# Every chunk but the continuation should have been written in place, i.e.
|
||||
# straight into the returned accumulator, with no intermediate copy.
|
||||
wrote_in_place = [
|
||||
out is not None and out.data_ptr() == fused_out[chunk.token_slice].data_ptr()
|
||||
for out, chunk in zip(
|
||||
backend_fused.out_destinations,
|
||||
prefill_fused.chunked_context.chunks,
|
||||
strict=True,
|
||||
)
|
||||
]
|
||||
continuations = [c.is_continuation for c in prefill_fused.chunked_context.chunks]
|
||||
assert any(continuations), "the batch must exercise a continuation chunk"
|
||||
if honors_out:
|
||||
assert wrote_in_place == [not c for c in continuations]
|
||||
else:
|
||||
assert not any(wrote_in_place)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_fused_context_rejects_an_unquantized_query() -> None:
|
||||
"""The fp8 query is produced by the new-token epilogue, not re-cast here."""
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda")
|
||||
fp8 = current_platform.fp8_dtype()
|
||||
prefill, num_blocks = _build_prefill_metadata(
|
||||
device, fp8, fp8, _RecordingPrefillBackend()
|
||||
)
|
||||
layer = _FusedLayer(
|
||||
_KVBProj(device, weight_dtype=fp8),
|
||||
torch.zeros((num_blocks, _BLOCK_SIZE, _ENTRY), device=device, dtype=fp8),
|
||||
"fp8",
|
||||
torch.ones(1, dtype=torch.float32, device=device),
|
||||
)
|
||||
q = torch.zeros(
|
||||
(sum(_QUERY_LENS), _NUM_HEADS, _QK_NOPE + _QK_ROPE),
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
with pytest.raises(AssertionError, match="new-token epilogue"):
|
||||
layer._compute_prefill_context(q, SimpleNamespace(prefill=prefill))
|
||||
@@ -24,7 +24,7 @@ class CustomMLAPrefillBackend(MLAPrefillBackend):
|
||||
def run_prefill_new_tokens(self, q, k, v, return_softmax_lse):
|
||||
raise NotImplementedError
|
||||
|
||||
def run_prefill_context_chunk(self, chunk, q, k, v):
|
||||
def run_prefill_context_chunk(self, chunk, q, k, v, out=None):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
@@ -112,7 +112,7 @@ def test_register_custom_backend_as_decorator():
|
||||
def run_prefill_new_tokens(self, q, k, v, return_softmax_lse):
|
||||
raise NotImplementedError
|
||||
|
||||
def run_prefill_context_chunk(self, chunk, q, k, v):
|
||||
def run_prefill_context_chunk(self, chunk, q, k, v, out=None):
|
||||
raise NotImplementedError
|
||||
|
||||
assert MLAPrefillBackendEnum.CUSTOM.is_overridden()
|
||||
|
||||
@@ -2414,6 +2414,22 @@ def reorg_kvcache(
|
||||
return reorganized_kv_c_normed, reorganized_k_pe
|
||||
|
||||
|
||||
def neutralize_empty_context_partials(
|
||||
chunked_context: "MLACommonPrefillMetadata.ChunkedContextMetadata",
|
||||
output: torch.Tensor,
|
||||
output_lse: torch.Tensor,
|
||||
) -> None:
|
||||
"""Neutralize the partial of every prefill that no chunk covers.
|
||||
|
||||
A prefill without context is never gathered, so nothing would write its rows;
|
||||
a zero output with an ``-inf`` lse carries no weight into the final merge
|
||||
against the suffix partial.
|
||||
"""
|
||||
for token_slice in chunked_context.empty_token_slices:
|
||||
output[token_slice].zero_()
|
||||
output_lse[:, token_slice].fill_(float("-inf"))
|
||||
|
||||
|
||||
def init_mla_context_partial(
|
||||
chunked_context: "MLACommonPrefillMetadata.ChunkedContextMetadata",
|
||||
attn_output: torch.Tensor,
|
||||
@@ -2423,7 +2439,9 @@ def init_mla_context_partial(
|
||||
"""Allocate the running context partial over all prefill tokens.
|
||||
|
||||
Laid out like the chunk partials so the final whole-batch merge against the
|
||||
suffix partial sees matching head strides.
|
||||
suffix partial sees matching head strides. Callers whose backend honors an
|
||||
``out`` tensor already know that layout and can allocate directly, pairing it
|
||||
with ``neutralize_empty_context_partials``.
|
||||
"""
|
||||
output = torch.empty(
|
||||
(num_tokens, *attn_output.shape[1:]),
|
||||
@@ -2435,10 +2453,7 @@ def init_mla_context_partial(
|
||||
dtype=attn_softmax_lse.dtype,
|
||||
device=attn_softmax_lse.device,
|
||||
)
|
||||
# No chunk covers a prefill without context, so neutralize its partial.
|
||||
for token_slice in chunked_context.empty_token_slices:
|
||||
output[token_slice].zero_()
|
||||
output_lse[:, token_slice].fill_(float("-inf"))
|
||||
neutralize_empty_context_partials(chunked_context, output, output_lse)
|
||||
return output, output_lse
|
||||
|
||||
|
||||
@@ -2448,16 +2463,26 @@ def accumulate_mla_context_chunk(
|
||||
attn_softmax_lse: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
output_lse: torch.Tensor,
|
||||
output_written: bool = False,
|
||||
) -> None:
|
||||
"""Fold one chunk's partial into the running context partial.
|
||||
|
||||
Only the first request may be a continuation; its tokens are merged and the
|
||||
remaining token range is initialized.
|
||||
|
||||
Args:
|
||||
output_written: The chunk's attention output already landed in
|
||||
``output[chunk.token_slice]`` because the backend was handed it as
|
||||
``out``, leaving only the lse to fold. Invalid for a continuation
|
||||
chunk, whose leading tokens must be merged rather than overwritten.
|
||||
"""
|
||||
token_start = chunk.token_slice.start
|
||||
token_end = chunk.token_slice.stop
|
||||
init_start = token_start
|
||||
if chunk.is_continuation:
|
||||
assert not output_written, (
|
||||
"a continuation chunk must not write over the partial it merges with"
|
||||
)
|
||||
init_start = chunk.continuation_token_end
|
||||
num_merged = init_start - token_start
|
||||
merge_attn_states(
|
||||
@@ -2470,7 +2495,8 @@ def accumulate_mla_context_chunk(
|
||||
)
|
||||
if init_start < token_end:
|
||||
written = init_start - token_start
|
||||
output[init_start:token_end].copy_(attn_output[written:])
|
||||
if not output_written:
|
||||
output[init_start:token_end].copy_(attn_output[written:])
|
||||
output_lse[:, init_start:token_end].copy_(attn_softmax_lse[:, written:])
|
||||
|
||||
|
||||
|
||||
@@ -8,8 +8,9 @@ This is a self-contained MLA layer that owns the full attention path:
|
||||
-> fused pre-attention ops (fused_qkv_a_proj / norms / q_b_proj)
|
||||
-> explicit prefill / decode split
|
||||
prefill: fused key-concat + cache-insert kernel -> run_prefill_new_tokens
|
||||
(+ chunked-context merge); dispatched by cache dtype
|
||||
(bf16 / plain fp8 / fp8_ds_mla)
|
||||
(+ chunked-context merge, whose per-chunk gather -> kv_b_proj
|
||||
-> fused K/V pack loop this layer owns); dispatched by cache
|
||||
dtype (bf16 / plain fp8 / fp8_ds_mla)
|
||||
decode : W_UK absorb (BMM1) -> fused q-concat + cache-insert kernel
|
||||
-> impl.forward_mqa -> W_UV up-proj (MQA)
|
||||
-> optional output gate
|
||||
@@ -34,6 +35,7 @@ from typing import TYPE_CHECKING, cast
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
|
||||
from vllm.config import (
|
||||
CacheConfig,
|
||||
@@ -48,6 +50,12 @@ from vllm.model_executor.layers.attention.attention import (
|
||||
set_default_quant_scales,
|
||||
should_load_quant_weights,
|
||||
)
|
||||
from vllm.model_executor.layers.attention.mla_attention import (
|
||||
_get_kv_b_proj_input_dtype,
|
||||
accumulate_mla_context_chunk,
|
||||
init_mla_context_partial,
|
||||
neutralize_empty_context_partials,
|
||||
)
|
||||
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.model_executor.layers.linear import (
|
||||
@@ -67,6 +75,8 @@ from vllm.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import (
|
||||
fused_mla_decode_q_concat_kv_cache_insert,
|
||||
fused_mla_key_concat_ds_mla_insert,
|
||||
fused_mla_key_concat_kv_cache_insert,
|
||||
fused_mla_kv_concat,
|
||||
fused_mla_kv_concat_quant_fp8,
|
||||
fused_mla_qkv_quant_kv_cache_fp8_insert,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
@@ -696,6 +706,181 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
|
||||
cos_sin_cache=cos_sin_cache,
|
||||
)
|
||||
|
||||
def _compute_prefill_context(
|
||||
self,
|
||||
q: torch.Tensor,
|
||||
attn_metadata: "MLACommonMetadata",
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Chunked-context prefill, K3-fused. Replaces the impl's version.
|
||||
|
||||
Per chunk the impl gathers the paged latent, up-projects it, then casts
|
||||
and concatenates K (and casts V) in two or three more launches. Here that
|
||||
tail is one fused kernel per chunk -- ``fused_mla_kv_concat`` for a bf16
|
||||
query, ``fused_mla_kv_concat_quant_fp8`` when the query is fp8 -- reading
|
||||
the strided ``kv_b_proj`` output in place and writing a contiguous key, so
|
||||
only the gather and ``kv_b_proj`` remain.
|
||||
|
||||
The impl's query cast is gone as well: ``q`` already carries
|
||||
``prefill.q_data_type`` because the new-token epilogue quantized it. The
|
||||
gathered latent still gets the impl's cast to whatever ``kv_b_proj``
|
||||
consumes -- free (a no-op ``.to``) for a checkpoint whose ``kv_b_proj``
|
||||
takes the fp8 latent directly, and required for a bf16 one, which is
|
||||
what a stock K3 checkpoint carries. Its output is bf16 either way.
|
||||
|
||||
The gathered ``k_pe`` is likewise used as-is (fp8 for a plain fp8 cache)
|
||||
and needs no RoPE: it was rotated on the way in.
|
||||
|
||||
Chunk partials are written straight into the accumulating context partial
|
||||
when the prefill backend honors ``out``, so only the (64x smaller) lse is
|
||||
copied per chunk.
|
||||
|
||||
Decode context parallelism keeps using
|
||||
``impl._context_parallel_compute_prefill_context``; its extra allgather
|
||||
and reorg are not fused here.
|
||||
"""
|
||||
prefill = attn_metadata.prefill
|
||||
assert prefill is not None
|
||||
prefill_backend = prefill.prefill_backend
|
||||
assert prefill_backend is not None
|
||||
chunked_context = prefill.chunked_context
|
||||
assert chunked_context is not None
|
||||
assert q.dtype == prefill.q_data_type, (
|
||||
"Kimi-K3 chunked context expects the new-token epilogue to have "
|
||||
f"produced a {prefill.q_data_type} query; got {q.dtype}."
|
||||
)
|
||||
|
||||
fp8_prefill = q.dtype == current_platform.fp8_dtype()
|
||||
workspace = chunked_context.workspace
|
||||
kv_cache = self._attn_read_kv_cache()
|
||||
kv_b_proj_input_dtype = _get_kv_b_proj_input_dtype(self.kv_b_proj, fp8_prefill)
|
||||
|
||||
def run_chunk(
|
||||
chunk, out: torch.Tensor | None = None
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
self._gather_context_latent(chunk, kv_cache, prefill, fp8_prefill)
|
||||
gathered = workspace[: chunk.num_context_tokens]
|
||||
kv_c_normed = gathered[..., : self.kv_lora_rank]
|
||||
if kv_b_proj_input_dtype is not None:
|
||||
kv_c_normed = kv_c_normed.to(kv_b_proj_input_dtype)
|
||||
kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
|
||||
-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim
|
||||
)
|
||||
k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
|
||||
k_pe = gathered[..., self.kv_lora_rank :]
|
||||
if fp8_prefill:
|
||||
k, v = fused_mla_kv_concat_quant_fp8(k_nope, k_pe, v)
|
||||
else:
|
||||
k = fused_mla_kv_concat(k_nope, k_pe)
|
||||
attn_output, attn_lse = prefill_backend.run_prefill_context_chunk(
|
||||
chunk=chunk, q=q[chunk.token_slice], k=k, v=v, out=out
|
||||
)
|
||||
assert out is None or attn_output.data_ptr() == out.data_ptr(), (
|
||||
f"{prefill_backend.get_name()} reports supports_out() but did not "
|
||||
"write the context chunk into the `out` it was given."
|
||||
)
|
||||
return attn_output, attn_lse
|
||||
|
||||
chunks = chunked_context.chunks
|
||||
if len(chunks) == 1 and not chunked_context.empty_token_slices:
|
||||
# One chunk covering every prefill token: its partial *is* the context
|
||||
# partial, so it needs neither an accumulator nor a copy.
|
||||
return run_chunk(chunks[0])
|
||||
|
||||
# A backend honoring `out` writes each chunk's partial straight into the
|
||||
# accumulator, so the per-chunk output copy disappears -- and because that
|
||||
# contract fixes the trailing shape, the accumulator can be sized before
|
||||
# any chunk runs. Otherwise the shape is only knowable from a real partial,
|
||||
# so the first chunk runs ahead of the loop and is copied in.
|
||||
writes_out = prefill_backend.supports_out()
|
||||
if writes_out:
|
||||
assert prefill.output_dtype is not None
|
||||
output = torch.empty(
|
||||
(q.shape[0], self.num_local_heads, self.v_head_dim),
|
||||
dtype=prefill.output_dtype,
|
||||
device=q.device,
|
||||
)
|
||||
output_lse = torch.empty(
|
||||
(self.num_local_heads, q.shape[0]),
|
||||
dtype=torch.float32,
|
||||
device=q.device,
|
||||
)
|
||||
neutralize_empty_context_partials(chunked_context, output, output_lse)
|
||||
else:
|
||||
attn_output, attn_lse = run_chunk(chunks[0])
|
||||
output, output_lse = init_mla_context_partial(
|
||||
chunked_context, attn_output, attn_lse, num_tokens=q.shape[0]
|
||||
)
|
||||
accumulate_mla_context_chunk(
|
||||
chunks[0], attn_output, attn_lse, output, output_lse
|
||||
)
|
||||
chunks = chunks[1:]
|
||||
|
||||
for chunk in chunks:
|
||||
# A continuation chunk's leading tokens have to be merged with the
|
||||
# partial already sitting there, so it cannot write in place.
|
||||
out = (
|
||||
output[chunk.token_slice]
|
||||
if writes_out and not chunk.is_continuation
|
||||
else None
|
||||
)
|
||||
attn_output, attn_lse = run_chunk(chunk, out=out)
|
||||
accumulate_mla_context_chunk(
|
||||
chunk,
|
||||
attn_output,
|
||||
attn_lse,
|
||||
output,
|
||||
output_lse,
|
||||
output_written=out is not None,
|
||||
)
|
||||
return output, output_lse
|
||||
|
||||
def _gather_context_latent(
|
||||
self,
|
||||
chunk,
|
||||
kv_cache: torch.Tensor,
|
||||
prefill,
|
||||
fp8_prefill: bool,
|
||||
) -> None:
|
||||
"""Gather one chunk's paged context latent into the workspace.
|
||||
|
||||
Dispatched exactly as in ``impl._compute_prefill_context``: an fp8 query
|
||||
reads the plain fp8 cache in its stored layout, anything else lands in the
|
||||
workspace as the model dtype.
|
||||
"""
|
||||
workspace = prefill.chunked_context.workspace
|
||||
toks = chunk.num_context_tokens
|
||||
block_table = prefill.block_table[chunk.request_slice]
|
||||
if self.kv_cache_dtype == "fp8_ds_mla":
|
||||
ops.cp_gather_and_upconvert_fp8_kv_cache(
|
||||
src_cache=kv_cache,
|
||||
dst=workspace[:toks],
|
||||
block_table=block_table,
|
||||
workspace_starts=chunk.cu_seq_lens,
|
||||
batch_size=chunk.num_requests,
|
||||
seq_starts=chunk.starts,
|
||||
)
|
||||
elif not fp8_prefill:
|
||||
ops.gather_and_maybe_dequant_cache(
|
||||
src_cache=kv_cache,
|
||||
dst=workspace,
|
||||
block_table=block_table,
|
||||
cu_seq_lens=chunk.cu_seq_lens,
|
||||
token_to_seq=chunk.token_to_seq,
|
||||
num_tokens=toks,
|
||||
kv_cache_dtype=self.kv_cache_dtype,
|
||||
scale=self._k_scale,
|
||||
seq_starts=chunk.starts,
|
||||
)
|
||||
else:
|
||||
ops.cp_gather_cache(
|
||||
src_cache=kv_cache,
|
||||
dst=workspace[:toks],
|
||||
block_table=block_table,
|
||||
cu_seq_lens=chunk.cu_seq_lens,
|
||||
batch_size=chunk.num_requests,
|
||||
seq_starts=chunk.starts,
|
||||
)
|
||||
|
||||
def _forward_prefill_fused(
|
||||
self,
|
||||
q: torch.Tensor,
|
||||
@@ -710,8 +895,9 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
|
||||
"""Prefill using the fused key-concat + cache-insert kernel.
|
||||
|
||||
Replaces ``_concat_k_nope_k_pe`` and the prefill cache write with one
|
||||
fused kernel launch, dispatched by cache dtype. The chunked context
|
||||
gather + online-softmax merge are delegated to the impl.
|
||||
fused kernel launch, dispatched by cache dtype. Chunked context runs
|
||||
through this layer's ``_compute_prefill_context``, except under DCP where
|
||||
it is delegated to the impl.
|
||||
|
||||
Supported configs (K3 fp8 policy):
|
||||
- bf16 cache -> bf16 prefill query
|
||||
@@ -812,8 +998,8 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
|
||||
)
|
||||
)
|
||||
else:
|
||||
context_output, context_lse = self.impl._compute_prefill_context( # type: ignore[attr-defined]
|
||||
q, self._attn_read_kv_cache(), attn_metadata, self._k_scale
|
||||
context_output, context_lse = self._compute_prefill_context(
|
||||
q, attn_metadata
|
||||
)
|
||||
suffix_output, suffix_lse = output_prefill
|
||||
out = out.view(-1, self.num_local_heads, self.v_head_dim)
|
||||
|
||||
@@ -12,6 +12,9 @@ mirror ``fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_{bf16,fp8}_insert``.
|
||||
- ``fused_mla_qkv_quant_kv_cache_fp8_insert`` (fp8): additionally quantize
|
||||
``q``/``k``/``v`` to E4M3 with ``q_scale`` / ``k_scale`` / ``v_scale`` (the
|
||||
cache shares ``k_scale``, as in ``concat_and_cache_mla``).
|
||||
- ``fused_mla_kv_concat`` / ``fused_mla_kv_concat_quant_fp8`` (chunked context):
|
||||
the same K concat (plus the fp8 K/V cast) without the cache insert, for context
|
||||
chunks whose latent was gathered back out of the paged cache.
|
||||
|
||||
The optional ``positions`` / ``cos_sin_cache`` pair enables GPT-J-style RoPE
|
||||
inside the epilogue. Omitting both keeps the K3 NoPE fast path. The kernels use
|
||||
@@ -156,6 +159,59 @@ def fused_mla_qkv_quant_kv_cache_fp8_insert(
|
||||
return q_fp8, k_fp8, v_fp8
|
||||
|
||||
|
||||
def _empty_full_key(
|
||||
k_nope: torch.Tensor, k_pe: torch.Tensor, dtype: torch.dtype
|
||||
) -> torch.Tensor:
|
||||
num_tokens, num_heads, qk_nope_head_dim = k_nope.shape
|
||||
return torch.empty(
|
||||
(num_tokens, num_heads, qk_nope_head_dim + k_pe.shape[1]),
|
||||
dtype=dtype,
|
||||
device=k_nope.device,
|
||||
)
|
||||
|
||||
|
||||
def fused_mla_kv_concat(
|
||||
k_nope: torch.Tensor, # [T, H, qk_nope_head_dim], may be strided
|
||||
k_pe: torch.Tensor, # [T, rope] or [T, 1, rope], same dtype as k_nope
|
||||
) -> torch.Tensor:
|
||||
"""Concat ``k = [k_nope | k_pe]`` into a contiguous key, in one launch.
|
||||
|
||||
The chunked-context counterpart of ``fused_mla_key_concat_kv_cache_insert``:
|
||||
no query, no cache insert and no RoPE (the gathered ``k_pe`` is already
|
||||
rotated). ``k_nope`` is a strided half of one ``kv_b_proj`` output and
|
||||
``k_pe`` a strided view of the gather workspace, so neither has to be made
|
||||
contiguous first.
|
||||
"""
|
||||
k_pe = k_pe.reshape(k_pe.shape[0], k_pe.shape[-1])
|
||||
k = _empty_full_key(k_nope, k_pe, k_nope.dtype)
|
||||
if k.shape[0]:
|
||||
torch.ops._C.fused_kimi_k3_mla_kv_concat(k_nope, k_pe, k)
|
||||
return k
|
||||
|
||||
|
||||
def fused_mla_kv_concat_quant_fp8(
|
||||
k_nope: torch.Tensor, # [T, H, qk_nope_head_dim], may be strided
|
||||
k_pe: torch.Tensor, # [T, rope] or [T, 1, rope], k_nope's dtype or fp8
|
||||
v: torch.Tensor, # [T, H, v_head_dim], may be strided
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""``fused_mla_kv_concat`` plus an fp8 cast of the key and of ``v``.
|
||||
|
||||
``k_pe`` may already be fp8: a plain fp8 cache is gathered without
|
||||
dequantizing, and those bytes are copied through as-is.
|
||||
|
||||
Returns contiguous ``(k_fp8, v_fp8)``.
|
||||
"""
|
||||
k_pe = k_pe.reshape(k_pe.shape[0], k_pe.shape[-1])
|
||||
fp8 = torch.float8_e4m3fn
|
||||
k_fp8 = _empty_full_key(k_nope, k_pe, fp8)
|
||||
v_fp8 = torch.empty(v.shape, dtype=fp8, device=v.device)
|
||||
if k_fp8.shape[0]:
|
||||
torch.ops._C.fused_kimi_k3_mla_kv_concat_quant_fp8(
|
||||
k_nope, k_pe, v, k_fp8, v_fp8
|
||||
)
|
||||
return k_fp8, v_fp8
|
||||
|
||||
|
||||
def fused_mla_decode_q_concat_kv_cache_insert(
|
||||
ql_nope: torch.Tensor, # [B, H, kv_lora_rank] (BMM1 output, absorbed q)
|
||||
q_pe: torch.Tensor, # [B, H, qk_rope_head_dim]
|
||||
|
||||
@@ -106,7 +106,12 @@ class AiterFlashAttnPrefillBackend(MLAPrefillBackend):
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
out: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
assert out is None, (
|
||||
"AiterFlashAttnPrefillBackend does not report supports_out(), so it "
|
||||
"is never given a context-chunk `out` to write into."
|
||||
)
|
||||
out, lse = self.flash_attn_varlen_func(
|
||||
q=q,
|
||||
k=k,
|
||||
|
||||
@@ -74,13 +74,16 @@ class MLAPrefillBackend(ABC):
|
||||
return False
|
||||
|
||||
def supports_out(self) -> bool:
|
||||
"""Whether `run_prefill_new_tokens` honors a caller-provided `out`
|
||||
tensor of shape `[num_tokens, num_heads, v_head_dim]`, writing the
|
||||
final result into it in place.
|
||||
"""Whether `run_prefill_new_tokens` and `run_prefill_context_chunk` honor
|
||||
a caller-provided `out` tensor of shape
|
||||
`[num_tokens, num_heads, v_head_dim]`, writing the result into it in place
|
||||
and returning it.
|
||||
|
||||
When True, callers may pass `out` and skip the post-hoc
|
||||
slice/flatten/copy. False for backends that ignore `out` or emit a
|
||||
padded (`qk_head_dim`) output. Overridden by backends that support it.
|
||||
slice/flatten/copy -- and, for context chunks, size the accumulating
|
||||
partial before running any chunk. False for backends that ignore `out` or
|
||||
emit a padded (`qk_head_dim`) output. Overridden by backends that support
|
||||
it.
|
||||
"""
|
||||
return False
|
||||
|
||||
@@ -179,5 +182,6 @@ class MLAPrefillBackend(ABC):
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
out: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
raise NotImplementedError
|
||||
|
||||
@@ -477,6 +477,7 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
out: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
return self._flash_attn_varlen_diff_headdims(
|
||||
q=q,
|
||||
@@ -489,6 +490,7 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
|
||||
softmax_scale=self.scale,
|
||||
causal=False, # Context is unmasked
|
||||
return_softmax_lse=True,
|
||||
out=out,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -228,11 +228,13 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
out: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
attn_out, lse = self._prefill_chunks[chunk.index].run(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
out=out,
|
||||
return_lse=True,
|
||||
)
|
||||
|
||||
|
||||
@@ -166,11 +166,13 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend):
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
out: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
from tokenspeed_mla import tokenspeed_mla_prefill
|
||||
|
||||
# See note in run_prefill_new_tokens — `v` is a split-view of `kv_nope`
|
||||
# in `_compute_prefill_context` and arrives non-contiguous.
|
||||
# in the generic `_compute_prefill_context` and arrives non-contiguous.
|
||||
# (K3's fused fp8 K/V pack already hands over a contiguous `v`.)
|
||||
v = v.contiguous()
|
||||
|
||||
attn_out, lse = tokenspeed_mla_prefill(
|
||||
@@ -187,6 +189,7 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend):
|
||||
cum_seq_lens_q=chunk.query_start_loc,
|
||||
max_seq_len_q=chunk.max_query_len,
|
||||
enable_pdl=False,
|
||||
out=out,
|
||||
)
|
||||
|
||||
# Convert from (q_len, num_heads) to (num_heads, q_len)
|
||||
|
||||
@@ -147,16 +147,18 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
out: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
from flashinfer.prefill import trtllm_ragged_attention_deepseek
|
||||
|
||||
out = torch.empty(
|
||||
q.shape[0],
|
||||
q.shape[1],
|
||||
v.shape[2],
|
||||
device=q.device,
|
||||
dtype=self._prefill_metadata.output_dtype,
|
||||
)
|
||||
if out is None:
|
||||
out = torch.empty(
|
||||
q.shape[0],
|
||||
q.shape[1],
|
||||
v.shape[2],
|
||||
device=q.device,
|
||||
dtype=self._prefill_metadata.output_dtype,
|
||||
)
|
||||
|
||||
attn_out, lse = trtllm_ragged_attention_deepseek(
|
||||
query=q,
|
||||
|
||||
Reference in New Issue
Block a user