A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 17:01 UTC
[8/n] Migrate merge_attn_states, mamba, sampler to torch stable ABI (continued) (#43361)
Signed-off-by: Mikayla Gawarecki <mikaylagawarecki@gmail.com> Signed-off-by: Chris Leonard <chleonar@redhat.com> Co-authored-by: Mikayla Gawarecki <mikaylagawarecki@gmail.com> Co-authored-by: Shengqi Chen <harry-chen@outlook.com>
Chris Leonard committed
May 27, 2026 at 12:35 UTC
284e6f543d462016fc80c055ccbf088832c63129
13 files changed
+432
-403
CMakeLists.txt
+5
-5
@@ -305,14 +305,10 @@ endif()
305
#
306
307
set(VLLM_EXT_SRC
308
- "csrc/mamba/mamba_ssm/selective_scan_fwd.cu"
308
"csrc/cache_kernels.cu"
309
"csrc/cache_kernels_fused.cu"
310
"csrc/attention/paged_attention_v1.cu"
311
"csrc/attention/paged_attention_v2.cu"
313
- "csrc/attention/merge_attn_states.cu"
314
- "csrc/sampler.cu"
315
- "csrc/topk.cu"
312
"csrc/cuda_view.cu"
313
"csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
314
"csrc/quantization/activation_kernels.cu"
@@ -633,7 +629,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
629
"csrc/libtorch_stable/fused_qknorm_rope_kernel.cu"
630
"csrc/libtorch_stable/layernorm_kernels.cu"
631
"csrc/libtorch_stable/layernorm_quant_kernels.cu"
636
- "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu")
632
+ "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
633
+ "csrc/libtorch_stable/attention/merge_attn_states.cu"
634
+ "csrc/libtorch_stable/sampler.cu"
635
+ "csrc/libtorch_stable/topk.cu"
636
+ "csrc/libtorch_stable/mamba/selective_scan_fwd.cu")
637
638
if(VLLM_GPU_LANG STREQUAL "CUDA")
639
list(APPEND VLLM_STABLE_EXT_SRC
csrc/libtorch_stable/attention/merge_attn_states.cu
renamed
+56
-49
@@ -1,14 +1,14 @@
1
#include <optional>
2
-#include <torch/all.h>
3
-#include <ATen/cuda/CUDAContext.h>
4
-#include <c10/cuda/CUDAGuard.h>
2
#include <algorithm>
3
#include <limits>
4
8
-#include "attention_dtypes.h"
9
-#include "attention_utils.cuh"
10
-#include "../quantization/w8a8/fp8/common.cuh"
5
+#include "../torch_utils.h"
6
#include "../dispatch_utils.h"
7
+#include <torch/headeronly/core/ScalarType.h>
8
+
9
+#include "../../attention/attention_dtypes.h"
10
+#include "../../attention/attention_utils.cuh"
11
+#include "../../quantization/w8a8/fp8/common.cuh"
12
13
namespace vllm {
14
@@ -196,17 +196,17 @@ __global__ void merge_attn_states_kernel(
196
// The following macro is used to dispatch the conversion function based on
197
// the output data type. The FN is a macro that calls a function with
198
// template<typename scalar_t>.
199
-#define DISPATCH_BY_SCALAR_DTYPE(scalar_dtype, fn) \
200
- { \
201
- if (scalar_dtype == at::ScalarType::Float) { \
202
- fn(float); \
203
- } else if (scalar_dtype == at::ScalarType::Half) { \
204
- fn(uint16_t); \
205
- } else if (scalar_dtype == at::ScalarType::BFloat16) { \
206
- fn(__nv_bfloat16); \
207
- } else { \
208
- TORCH_CHECK(false, "Unsupported data type of O: ", scalar_dtype); \
209
- } \
199
+#define DISPATCH_BY_SCALAR_DTYPE(scalar_dtype, fn) \
200
+ { \
201
+ if (scalar_dtype == torch::headeronly::ScalarType::Float) { \
202
+ fn(float); \
203
+ } else if (scalar_dtype == torch::headeronly::ScalarType::Half) { \
204
+ fn(uint16_t); \
205
+ } else if (scalar_dtype == torch::headeronly::ScalarType::BFloat16) { \
206
+ fn(__nv_bfloat16); \
207
+ } else { \
208
+ STD_TORCH_CHECK(false, "Unsupported data type of O: ", scalar_dtype); \
209
+ } \
210
}
211
212
#define LAUNCH_MERGE_ATTN_STATES(scalar_t, output_t, NUM_THREADS, \
@@ -245,11 +245,14 @@ __global__ void merge_attn_states_kernel(
245
*/
246
template <typename scalar_t>
247
void merge_attn_states_launcher(
248
- torch::Tensor& output, std::optional<torch::Tensor> output_lse,
249
- const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse,
250
- const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse,
248
+ torch::stable::Tensor& output,
249
+ std::optional<torch::stable::Tensor> output_lse,
250
+ const torch::stable::Tensor& prefix_output,
251
+ const torch::stable::Tensor& prefix_lse,
252
+ const torch::stable::Tensor& suffix_output,
253
+ const torch::stable::Tensor& suffix_lse,
254
const std::optional<int64_t> prefill_tokens_with_context,
252
- const std::optional<torch::Tensor>& output_scale) {
255
+ const std::optional<torch::stable::Tensor>& output_scale) {
256
constexpr uint NUM_THREADS = 128;
257
const uint num_tokens = output.size(0);
258
const uint num_heads = output.size(1);
@@ -258,23 +261,23 @@ void merge_attn_states_launcher(
261
const uint output_head_stride = output.stride(1);
262
// Thread mapping is based on input BF16 pack_size
263
const uint pack_size = 16 / sizeof(scalar_t);
261
- TORCH_CHECK(head_size % pack_size == 0,
262
- "headsize must be multiple of pack_size:", pack_size);
264
+ STD_TORCH_CHECK(head_size % pack_size == 0,
265
+ "headsize must be multiple of pack_size:", pack_size);
266
267
const uint prefix_num_tokens =
268
prefill_tokens_with_context.has_value()
269
? static_cast<uint>(prefill_tokens_with_context.value())
270
: num_tokens;
268
- TORCH_CHECK(prefix_num_tokens <= num_tokens,
269
- "prefix_num_tokens must be <= num_tokens");
271
+ STD_TORCH_CHECK(prefix_num_tokens <= num_tokens,
272
+ "prefix_num_tokens must be <= num_tokens");
273
274
float* output_lse_ptr = nullptr;
275
if (output_lse.has_value()) {
273
- output_lse_ptr = output_lse.value().data_ptr<float>();
276
+ output_lse_ptr = output_lse.value().mutable_data_ptr<float>();
277
}
278
float* output_scale_ptr = nullptr;
279
if (output_scale.has_value()) {
277
- output_scale_ptr = output_scale.value().data_ptr<float>();
280
+ output_scale_ptr = output_scale.value().mutable_data_ptr<float>();
281
}
282
// Process one pack elements per thread. for float, the
283
// pack_size is 4 for half/bf16, the pack_size is 8.
@@ -284,14 +287,15 @@ void merge_attn_states_launcher(
287
dim3 block(NUM_THREADS);
288
dim3 grid((total_threads + NUM_THREADS - 1) / NUM_THREADS);
289
287
- const c10::cuda::OptionalCUDAGuard device_guard(prefix_output.device());
288
- auto stream = at::cuda::getCurrentCUDAStream();
290
+ const torch::stable::accelerator::DeviceGuard device_guard(
291
+ prefix_output.get_device_index());
292
+ auto stream = get_current_cuda_stream();
293
294
if (output_scale.has_value()) {
295
// FP8 output path - dispatch on output FP8 type
292
- VLLM_DISPATCH_FP8_TYPES(output.scalar_type(), "merge_attn_states_fp8", [&] {
293
- LAUNCH_MERGE_ATTN_STATES(scalar_t, fp8_t, NUM_THREADS, true);
294
- });
296
+ VLLM_STABLE_DISPATCH_FP8_TYPES(
297
+ output.scalar_type(), "merge_attn_states_fp8",
298
+ [&] { LAUNCH_MERGE_ATTN_STATES(scalar_t, fp8_t, NUM_THREADS, true); });
299
} else {
300
// Original BF16/FP16/FP32 output path
301
LAUNCH_MERGE_ATTN_STATES(scalar_t, scalar_t, NUM_THREADS, false);
@@ -305,26 +309,29 @@ void merge_attn_states_launcher(
309
suffix_lse, prefill_tokens_with_context, output_scale); \
310
}
311
308
-void merge_attn_states(torch::Tensor& output,
309
- std::optional<torch::Tensor> output_lse,
310
- const torch::Tensor& prefix_output,
311
- const torch::Tensor& prefix_lse,
312
- const torch::Tensor& suffix_output,
313
- const torch::Tensor& suffix_lse,
314
- std::optional<int64_t> prefill_tokens_with_context,
315
- const std::optional<torch::Tensor>& output_scale) {
312
+void merge_attn_states(
313
+ torch::stable::Tensor& output,
314
+ std::optional<torch::stable::Tensor> output_lse,
315
+ const torch::stable::Tensor& prefix_output,
316
+ const torch::stable::Tensor& prefix_lse,
317
+ const torch::stable::Tensor& suffix_output,
318
+ const torch::stable::Tensor& suffix_lse,
319
+ const std::optional<int64_t> prefill_tokens_with_context,
320
+ const std::optional<torch::stable::Tensor>& output_scale) {
321
if (output_scale.has_value()) {
317
- TORCH_CHECK(output.scalar_type() == at::ScalarType::Float8_e4m3fn ||
318
- output.scalar_type() == at::ScalarType::Float8_e4m3fnuz,
319
- "output must be FP8 when output_scale is provided, got: ",
320
- output.scalar_type());
322
+ STD_TORCH_CHECK(
323
+ output.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
324
+ output.scalar_type() ==
325
+ torch::headeronly::ScalarType::Float8_e4m3fnuz,
326
+ "output must be FP8 when output_scale is provided, got: ",
327
+ output.scalar_type());
328
} else {
322
- TORCH_CHECK(output.scalar_type() == prefix_output.scalar_type(),
323
- "output dtype (", output.scalar_type(),
324
- ") must match prefix_output dtype (",
325
- prefix_output.scalar_type(), ") when output_scale is not set");
329
+ STD_TORCH_CHECK(
330
+ output.scalar_type() == prefix_output.scalar_type(), "output dtype (",
331
+ output.scalar_type(), ") must match prefix_output dtype (",
332
+ prefix_output.scalar_type(), ") when output_scale is not set");
333
}
334
// Always dispatch on prefix_output (input) dtype
328
- DISPATCH_BY_SCALAR_DTYPE(prefix_output.dtype(),
335
+ DISPATCH_BY_SCALAR_DTYPE(prefix_output.scalar_type(),
336
CALL_MERGE_ATTN_STATES_LAUNCHER);
337
}
csrc/libtorch_stable/mamba/selective_scan.h
renamed
+7
-4
@@ -12,6 +12,9 @@
12
#include <hip/hip_bf16.h>
13
#endif
14
#include <cuda_fp16.h>
15
+
16
+#include <torch/headeronly/util/Half.h>
17
+#include <torch/headeronly/util/BFloat16.h>
18
////////////////////////////////////////////////////////////////////////////////////////////////////
19
20
struct SSMParamsBase {
@@ -159,8 +162,8 @@ struct Converter{
162
};
163
164
template<int N>
162
-struct Converter<at::Half, N>{
163
- static inline __device__ void to_float(const at::Half (&src)[N], float (&dst)[N]) {
165
+struct Converter<torch::headeronly::Half, N>{
166
+ static inline __device__ void to_float(const torch::headeronly::Half (&src)[N], float (&dst)[N]) {
167
static_assert(N % 2 == 0);
168
auto &src2 = reinterpret_cast<const half2 (&)[N / 2]>(src);
169
auto &dst2 = reinterpret_cast<float2 (&)[N / 2]>(dst);
@@ -171,8 +174,8 @@ struct Converter<at::Half, N>{
174
175
#if __CUDA_ARCH__ >= 800
176
template<int N>
174
-struct Converter<at::BFloat16, N>{
175
- static inline __device__ void to_float(const at::BFloat16 (&src)[N], float (&dst)[N]) {
177
+struct Converter<torch::headeronly::BFloat16, N>{
178
+ static inline __device__ void to_float(const torch::headeronly::BFloat16 (&src)[N], float (&dst)[N]) {
179
static_assert(N % 2 == 0);
180
auto &src2 = reinterpret_cast<const nv_bfloat162 (&)[N / 2]>(src);
181
auto &dst2 = reinterpret_cast<float2 (&)[N / 2]>(dst);
csrc/libtorch_stable/mamba/selective_scan_fwd.cu
renamed
+102
-111
@@ -1,18 +1,9 @@
1
// clang-format off
2
// adapted from https://github.com/state-spaces/mamba/blob/main/csrc/selective_scan/selective_scan_fwd_kernel.cuh
3
-#include <torch/all.h>
4
-#include <ATen/cuda/CUDAContext.h>
5
-#include <c10/cuda/CUDAGuard.h>
3
+#include "../torch_utils.h"
4
+#include <torch/csrc/stable/macros.h>
5
#include "selective_scan.h"
6
8
-#include <c10/util/BFloat16.h>
9
-#include <c10/util/Half.h>
10
-#ifdef USE_ROCM
11
- #include <c10/hip/HIPException.h> // For C10_HIP_CHECK and C10_HIP_KERNEL_LAUNCH_CHECK
12
-#else
13
- #include <c10/cuda/CUDAException.h> // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
14
-#endif
15
-
7
#ifndef USE_ROCM
8
#include <cub/block/block_load.cuh>
9
#include <cub/block/block_store.cuh>
@@ -416,15 +407,15 @@ void selective_scan_fwd_launch(SSMParamsBase ¶ms, cudaStream_t stream) {
407
auto kernel = &selective_scan_fwd_kernel<Ktraits>;
408
if (kSmemSize >= 48 * 1024) {
409
#ifdef USE_ROCM
419
- C10_HIP_CHECK(hipFuncSetAttribute(
410
+ STD_CUDA_CHECK(hipFuncSetAttribute(
411
reinterpret_cast<const void*>(kernel), hipFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
412
#else
422
- C10_CUDA_CHECK(cudaFuncSetAttribute(
413
+ STD_CUDA_CHECK(cudaFuncSetAttribute(
414
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
415
#endif
416
}
417
kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
427
- C10_CUDA_KERNEL_LAUNCH_CHECK();
418
+ STD_CUDA_KERNEL_LAUNCH_CHECK();
419
});
420
});
421
});
@@ -462,46 +453,46 @@ void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream) {
453
#endif
454
}
455
465
-template void selective_scan_fwd_cuda<at::BFloat16, float, at::BFloat16>(SSMParamsBase ¶ms, cudaStream_t stream);
466
-template void selective_scan_fwd_cuda<at::BFloat16, float, float>(SSMParamsBase ¶ms, cudaStream_t stream);
467
-template void selective_scan_fwd_cuda<at::Half, float, at::Half>(SSMParamsBase ¶ms, cudaStream_t stream);
468
-template void selective_scan_fwd_cuda<at::Half, float, float>(SSMParamsBase ¶ms, cudaStream_t stream);
456
+template void selective_scan_fwd_cuda<torch::headeronly::BFloat16, float, torch::headeronly::BFloat16>(SSMParamsBase ¶ms, cudaStream_t stream);
457
+template void selective_scan_fwd_cuda<torch::headeronly::BFloat16, float, float>(SSMParamsBase ¶ms, cudaStream_t stream);
458
+template void selective_scan_fwd_cuda<torch::headeronly::Half, float, torch::headeronly::Half>(SSMParamsBase ¶ms, cudaStream_t stream);
459
+template void selective_scan_fwd_cuda<torch::headeronly::Half, float, float>(SSMParamsBase ¶ms, cudaStream_t stream);
460
template void selective_scan_fwd_cuda<float, float, float>(SSMParamsBase ¶ms, cudaStream_t stream);
461
471
-#define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
462
+#define CHECK_SHAPE(x, ...) STD_TORCH_CHECK(x.sizes().equals(torch::headeronly::IntHeaderOnlyArrayRef({__VA_ARGS__})), #x " must have shape (" #__VA_ARGS__ ")")
463
464
#define DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(ITYPE, STYPE, NAME, ...) \
474
- if (ITYPE == at::ScalarType::Half) { \
475
- using input_t = at::Half; \
465
+ if (ITYPE == torch::headeronly::ScalarType::Half) { \
466
+ using input_t = torch::headeronly::Half; \
467
using weight_t = float; \
477
- if (STYPE == at::ScalarType::Half) { \
478
- using state_t = at::Half; \
468
+ if (STYPE == torch::headeronly::ScalarType::Half) { \
469
+ using state_t = torch::headeronly::Half; \
470
__VA_ARGS__(); \
480
- } else if (STYPE == at::ScalarType::Float) { \
471
+ } else if (STYPE == torch::headeronly::ScalarType::Float) { \
472
using state_t = float; \
473
__VA_ARGS__(); \
474
} else { \
484
- AT_ERROR(#NAME, " not implemented for state type '", toString(STYPE), "'"); \
475
+ STD_TORCH_CHECK(false, #NAME " not implemented for state type '", STYPE, "'"); \
476
} \
486
- } else if (ITYPE == at::ScalarType::BFloat16) { \
487
- using input_t = at::BFloat16; \
477
+ } else if (ITYPE == torch::headeronly::ScalarType::BFloat16) { \
478
+ using input_t = torch::headeronly::BFloat16; \
479
using weight_t = float; \
489
- if (STYPE == at::ScalarType::BFloat16) { \
490
- using state_t = at::BFloat16; \
480
+ if (STYPE == torch::headeronly::ScalarType::BFloat16) { \
481
+ using state_t = torch::headeronly::BFloat16; \
482
__VA_ARGS__(); \
492
- } else if (STYPE == at::ScalarType::Float) { \
483
+ } else if (STYPE == torch::headeronly::ScalarType::Float) { \
484
using state_t = float; \
485
__VA_ARGS__(); \
486
} else { \
496
- AT_ERROR(#NAME, " not implemented for state type '", toString(STYPE), "'"); \
487
+ STD_TORCH_CHECK(false, #NAME " not implemented for state type '", STYPE, "'"); \
488
} \
498
- } else if (ITYPE == at::ScalarType::Float) { \
489
+ } else if (ITYPE == torch::headeronly::ScalarType::Float) { \
490
using input_t = float; \
491
using weight_t = float; \
492
using state_t = float; \
493
__VA_ARGS__(); \
494
} else { \
504
- AT_ERROR(#NAME, " not implemented for input type '", toString(ITYPE), "'"); \
495
+ STD_TORCH_CHECK(false, #NAME " not implemented for input type '", ITYPE, "'"); \
496
}
497
498
@@ -518,30 +509,30 @@ void set_ssm_params_fwd(SSMParamsBase ¶ms,
509
const bool is_variable_B,
510
const bool is_variable_C,
511
// device pointers
521
- const torch::Tensor u,
522
- const torch::Tensor delta,
523
- const torch::Tensor A,
524
- const torch::Tensor B,
525
- const torch::Tensor C,
526
- const torch::Tensor out,
527
- const torch::Tensor z,
528
- const torch::Tensor out_z,
529
- const std::optional<at::Tensor>& D,
530
- const std::optional<at::Tensor>& delta_bias,
531
- const torch::Tensor ssm_states,
512
+ const torch::stable::Tensor u,
513
+ const torch::stable::Tensor delta,
514
+ const torch::stable::Tensor A,
515
+ const torch::stable::Tensor B,
516
+ const torch::stable::Tensor C,
517
+ const torch::stable::Tensor out,
518
+ const torch::stable::Tensor z,
519
+ const torch::stable::Tensor out_z,
520
+ const std::optional<torch::stable::Tensor>& D,
521
+ const std::optional<torch::stable::Tensor>& delta_bias,
522
+ const torch::stable::Tensor ssm_states,
523
bool has_z,
524
bool delta_softplus,
534
- const std::optional<at::Tensor>& query_start_loc,
535
- const std::optional<at::Tensor>& cache_indices,
536
- const std::optional<at::Tensor>& has_initial_state,
525
+ const std::optional<torch::stable::Tensor>& query_start_loc,
526
+ const std::optional<torch::stable::Tensor>& cache_indices,
527
+ const std::optional<torch::stable::Tensor>& has_initial_state,
528
bool varlen,
529
int64_t null_block_id,
530
int64_t block_size,
540
- const std::optional<torch::Tensor> &block_idx_first_scheduled_token,
541
- const std::optional<torch::Tensor> &block_idx_last_scheduled_token,
542
- const std::optional<torch::Tensor> &initial_state_idx,
543
- const std::optional<torch::Tensor> &cu_chunk_seqlen,
544
- const std::optional<torch::Tensor> &last_chunk_indices) {
531
+ const std::optional<torch::stable::Tensor> &block_idx_first_scheduled_token,
532
+ const std::optional<torch::stable::Tensor> &block_idx_last_scheduled_token,
533
+ const std::optional<torch::stable::Tensor> &initial_state_idx,
534
+ const std::optional<torch::stable::Tensor> &cu_chunk_seqlen,
535
+ const std::optional<torch::stable::Tensor> &last_chunk_indices) {
536
537
// Reset the parameters
538
memset(¶ms, 0, sizeof(params));
@@ -654,45 +645,45 @@ void set_ssm_params_fwd(SSMParamsBase ¶ms,
645
}
646
}
647
657
-void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
658
- const torch::Tensor &A, const torch::Tensor &B, const torch::Tensor &C,
659
- const std::optional<torch::Tensor> &D_,
660
- const std::optional<torch::Tensor> &z_,
661
- const std::optional<torch::Tensor> &delta_bias_,
648
+void selective_scan_fwd(const torch::stable::Tensor &u, const torch::stable::Tensor &delta,
649
+ const torch::stable::Tensor &A, const torch::stable::Tensor &B, const torch::stable::Tensor &C,
650
+ const std::optional<torch::stable::Tensor> &D_,
651
+ const std::optional<torch::stable::Tensor> &z_,
652
+ const std::optional<torch::stable::Tensor> &delta_bias_,
653
bool delta_softplus,
663
- const std::optional<torch::Tensor> &query_start_loc,
664
- const std::optional<torch::Tensor> &cache_indices,
665
- const std::optional<torch::Tensor> &has_initial_state,
666
- const torch::Tensor &ssm_states,
654
+ const std::optional<torch::stable::Tensor> &query_start_loc,
655
+ const std::optional<torch::stable::Tensor> &cache_indices,
656
+ const std::optional<torch::stable::Tensor> &has_initial_state,
657
+ const torch::stable::Tensor &ssm_states,
658
// used to identify padding entries if cache_indices provided
659
// in case of padding, the kernel will return early
660
int64_t null_block_id,
661
int64_t block_size,
671
- const std::optional<torch::Tensor> &block_idx_first_scheduled_token,
672
- const std::optional<torch::Tensor> &block_idx_last_scheduled_token,
673
- const std::optional<torch::Tensor> &initial_state_idx,
674
- const std::optional<torch::Tensor> &cu_chunk_seqlen,
675
- const std::optional<torch::Tensor> &last_chunk_indices) {
662
+ const std::optional<torch::stable::Tensor> &block_idx_first_scheduled_token,
663
+ const std::optional<torch::stable::Tensor> &block_idx_last_scheduled_token,
664
+ const std::optional<torch::stable::Tensor> &initial_state_idx,
665
+ const std::optional<torch::stable::Tensor> &cu_chunk_seqlen,
666
+ const std::optional<torch::stable::Tensor> &last_chunk_indices) {
667
auto input_type = u.scalar_type();
668
auto weight_type = A.scalar_type();
678
- TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
679
- TORCH_CHECK(weight_type == at::ScalarType::Float);
669
+ STD_TORCH_CHECK(input_type == torch::headeronly::ScalarType::Float || input_type == torch::headeronly::ScalarType::Half || input_type == torch::headeronly::ScalarType::BFloat16);
670
+ STD_TORCH_CHECK(weight_type == torch::headeronly::ScalarType::Float);
671
672
const bool is_variable_B = B.dim() >= 3;
673
const bool is_variable_C = C.dim() >= 3;
674
684
- TORCH_CHECK(delta.scalar_type() == input_type);
685
- TORCH_CHECK(B.scalar_type() == (!is_variable_B ? weight_type : input_type));
686
- TORCH_CHECK(C.scalar_type() == (!is_variable_C ? weight_type : input_type));
675
+ STD_TORCH_CHECK(delta.scalar_type() == input_type);
676
+ STD_TORCH_CHECK(B.scalar_type() == (!is_variable_B ? weight_type : input_type));
677
+ STD_TORCH_CHECK(C.scalar_type() == (!is_variable_C ? weight_type : input_type));
678
688
- TORCH_CHECK(u.is_cuda());
689
- TORCH_CHECK(delta.is_cuda());
690
- TORCH_CHECK(A.is_cuda());
691
- TORCH_CHECK(B.is_cuda());
692
- TORCH_CHECK(C.is_cuda());
679
+ STD_TORCH_CHECK(u.is_cuda());
680
+ STD_TORCH_CHECK(delta.is_cuda());
681
+ STD_TORCH_CHECK(A.is_cuda());
682
+ STD_TORCH_CHECK(B.is_cuda());
683
+ STD_TORCH_CHECK(C.is_cuda());
684
694
- TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1);
695
- TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1);
685
+ STD_TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1);
686
+ STD_TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1);
687
688
const auto sizes = u.sizes();
689
const bool varlen = query_start_loc.has_value();
@@ -702,7 +693,7 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
693
const int dstate = A.size(1);
694
const int n_groups = varlen ? B.size(0) : B.size(1);
695
705
- TORCH_CHECK(dstate <= 256, "selective_scan only supports state dimension <= 256");
696
+ STD_TORCH_CHECK(dstate <= 256, "selective_scan only supports state dimension <= 256");
697
698
if (varlen) {
699
CHECK_SHAPE(u, dim, seqlen);
@@ -712,94 +703,94 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
703
CHECK_SHAPE(delta, batch_size, dim, seqlen);
704
}
705
CHECK_SHAPE(A, dim, dstate);
715
- TORCH_CHECK(is_variable_B, "is_variable_B = False is disabled in favor of reduced binary size")
706
+ STD_TORCH_CHECK(is_variable_B, "is_variable_B = False is disabled in favor of reduced binary size");
707
if (varlen) {
708
CHECK_SHAPE(B, n_groups, dstate, seqlen);
709
} else {
719
- CHECK_SHAPE(B, batch_size, n_groups, dstate, seqlen);
710
+ CHECK_SHAPE(B, batch_size, n_groups, dstate, seqlen);
711
}
721
- TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1);
712
+ STD_TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1);
713
723
- TORCH_CHECK(is_variable_C, "is_variable_C = False is disabled in favor of reduced binary size")
714
+ STD_TORCH_CHECK(is_variable_C, "is_variable_C = False is disabled in favor of reduced binary size");
715
if (varlen) {
716
CHECK_SHAPE(C, n_groups, dstate, seqlen);
717
} else {
727
- CHECK_SHAPE(C, batch_size, n_groups, dstate, seqlen);
718
+ CHECK_SHAPE(C, batch_size, n_groups, dstate, seqlen);
719
}
729
- TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1);
720
+ STD_TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1);
721
722
if (D_.has_value()) {
723
auto D = D_.value();
733
- TORCH_CHECK(D.scalar_type() == at::ScalarType::Float);
734
- TORCH_CHECK(D.is_cuda());
735
- TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1);
724
+ STD_TORCH_CHECK(D.scalar_type() == torch::headeronly::ScalarType::Float);
725
+ STD_TORCH_CHECK(D.is_cuda());
726
+ STD_TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1);
727
CHECK_SHAPE(D, dim);
728
}
729
730
if (delta_bias_.has_value()) {
731
auto delta_bias = delta_bias_.value();
741
- TORCH_CHECK(delta_bias.scalar_type() == at::ScalarType::Float);
742
- TORCH_CHECK(delta_bias.is_cuda());
743
- TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1);
732
+ STD_TORCH_CHECK(delta_bias.scalar_type() == torch::headeronly::ScalarType::Float);
733
+ STD_TORCH_CHECK(delta_bias.is_cuda());
734
+ STD_TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1);
735
CHECK_SHAPE(delta_bias, dim);
736
}
737
738
739
if (has_initial_state.has_value()) {
740
auto has_initial_state_ = has_initial_state.value();
750
- TORCH_CHECK(has_initial_state_.scalar_type() == at::ScalarType::Bool);
751
- TORCH_CHECK(has_initial_state_.is_cuda());
741
+ STD_TORCH_CHECK(has_initial_state_.scalar_type() == torch::headeronly::ScalarType::Bool);
742
+ STD_TORCH_CHECK(has_initial_state_.is_cuda());
743
CHECK_SHAPE(has_initial_state_, batch_size);
744
}
745
746
747
if (query_start_loc.has_value()) {
748
auto query_start_loc_ = query_start_loc.value();
758
- TORCH_CHECK(query_start_loc_.scalar_type() == at::ScalarType::Int);
759
- TORCH_CHECK(query_start_loc_.is_cuda());
749
+ STD_TORCH_CHECK(query_start_loc_.scalar_type() == torch::headeronly::ScalarType::Int);
750
+ STD_TORCH_CHECK(query_start_loc_.is_cuda());
751
}
752
753
754
if (cache_indices.has_value()) {
755
auto cache_indices_ = cache_indices.value();
765
- TORCH_CHECK(cache_indices_.scalar_type() == at::ScalarType::Int);
766
- TORCH_CHECK(cache_indices_.is_cuda());
756
+ STD_TORCH_CHECK(cache_indices_.scalar_type() == torch::headeronly::ScalarType::Int);
757
+ STD_TORCH_CHECK(cache_indices_.is_cuda());
758
759
// cache_indices can be either 1D (batch_size,) for non-APC mode
760
// or 2D (batch_size, max_positions) for APC mode
761
const bool is_apc_mode = block_idx_first_scheduled_token.has_value();
762
if (is_apc_mode) {
772
- TORCH_CHECK(cache_indices_.dim() == 2, "cache_indices must be 2D for APC mode");
773
- TORCH_CHECK(cache_indices_.size(0) == batch_size, "cache_indices first dimension must match batch_size");
763
+ STD_TORCH_CHECK(cache_indices_.dim() == 2, "cache_indices must be 2D for APC mode");
764
+ STD_TORCH_CHECK(cache_indices_.size(0) == batch_size, "cache_indices first dimension must match batch_size");
765
} else {
766
CHECK_SHAPE(cache_indices_, batch_size);
767
}
768
}
778
-
769
780
- at::Tensor z, out_z;
770
+
771
+ torch::stable::Tensor z, out_z;
772
const bool has_z = z_.has_value();
773
if (has_z) {
774
z = z_.value();
784
- TORCH_CHECK(z.scalar_type() == input_type);
785
- TORCH_CHECK(z.is_cuda());
786
- TORCH_CHECK(z.stride(-1) == 1 || z.size(-1) == 1);
775
+ STD_TORCH_CHECK(z.scalar_type() == input_type);
776
+ STD_TORCH_CHECK(z.is_cuda());
777
+ STD_TORCH_CHECK(z.stride(-1) == 1 || z.size(-1) == 1);
778
if (varlen){
779
CHECK_SHAPE(z, dim, seqlen);
780
} else {
781
CHECK_SHAPE(z, batch_size, dim, seqlen);
782
}
792
-
783
+
784
out_z = z;
785
}
786
787
// Right now u has BHL layout and delta has HBL layout, and we want out to have HBL layout
797
- at::Tensor out = delta;
788
+ torch::stable::Tensor out = delta;
789
// ssm_states can now be either the same as input_type or float32
790
auto state_type = ssm_states.scalar_type();
800
- TORCH_CHECK(state_type == input_type || state_type == at::ScalarType::Float);
801
- TORCH_CHECK(ssm_states.is_cuda());
802
- TORCH_CHECK(ssm_states.stride(-1) == 1);
791
+ STD_TORCH_CHECK(state_type == input_type || state_type == torch::headeronly::ScalarType::Float);
792
+ STD_TORCH_CHECK(ssm_states.is_cuda());
793
+ STD_TORCH_CHECK(ssm_states.stride(-1) == 1);
794
795
SSMParamsBase params;
796
set_ssm_params_fwd(params, batch_size, dim, seqlen, dstate, n_groups, is_variable_B, is_variable_C,
@@ -823,8 +814,8 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
814
);
815
816
826
- const at::cuda::OptionalCUDAGuard device_guard(device_of(u));
827
- auto stream = at::cuda::getCurrentCUDAStream().stream();
817
+ const torch::stable::accelerator::DeviceGuard device_guard(u.get_device_index());
818
+ auto stream = get_current_cuda_stream();
819
DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), ssm_states.scalar_type(), "selective_scan_fwd", [&] {
820
selective_scan_fwd_cuda<input_t, weight_t, state_t>(params, stream);
821
});
csrc/libtorch_stable/mamba/static_switch.h
renamed
csrc/libtorch_stable/ops.h
+53
@@ -164,6 +164,17 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
164
165
#endif
166
167
+// Attention kernels (shared CUDA/ROCm)
168
+void merge_attn_states(
169
+ torch::stable::Tensor& output,
170
+ std::optional<torch::stable::Tensor> output_lse,
171
+ const torch::stable::Tensor& prefix_output,
172
+ const torch::stable::Tensor& prefix_lse,
173
+ const torch::stable::Tensor& suffix_output,
174
+ const torch::stable::Tensor& suffix_lse,
175
+ const std::optional<int64_t> prefill_tokens_with_context,
176
+ const std::optional<torch::stable::Tensor>& output_scale = std::nullopt);
177
+
178
torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x,
179
bool inplace);
180
@@ -220,6 +231,48 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q,
231
torch::stable::Tensor& position_ids,
232
int64_t forced_token_heads_per_warp);
233
234
+// Sampler kernels (shared CUDA/ROCm)
235
+void apply_repetition_penalties_(
236
+ torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
237
+ const torch::stable::Tensor& output_mask,
238
+ const torch::stable::Tensor& repetition_penalties);
239
+
240
+void top_k_per_row_prefill(const torch::stable::Tensor& logits,
241
+ const torch::stable::Tensor& rowStarts,
242
+ const torch::stable::Tensor& rowEnds,
243
+ torch::stable::Tensor& indices, int64_t numRows,
244
+ int64_t stride0, int64_t stride1, int64_t topK);
245
+
246
+void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n,
247
+ const torch::stable::Tensor& seqLens,
248
+ torch::stable::Tensor& indices, int64_t numRows,
249
+ int64_t stride0, int64_t stride1, int64_t topK);
250
+
251
+void persistent_topk(const torch::stable::Tensor& logits,
252
+ const torch::stable::Tensor& lengths,
253
+ torch::stable::Tensor& output,
254
+ torch::stable::Tensor& workspace, int64_t k,
255
+ int64_t max_seq_len);
256
+
257
+void selective_scan_fwd(
258
+ const torch::stable::Tensor& u, const torch::stable::Tensor& delta,
259
+ const torch::stable::Tensor& A, const torch::stable::Tensor& B,
260
+ const torch::stable::Tensor& C,
261
+ const std::optional<torch::stable::Tensor>& D_,
262
+ const std::optional<torch::stable::Tensor>& z_,
263
+ const std::optional<torch::stable::Tensor>& delta_bias_,
264
+ bool delta_softplus,
265
+ const std::optional<torch::stable::Tensor>& query_start_loc,
266
+ const std::optional<torch::stable::Tensor>& cache_indices,
267
+ const std::optional<torch::stable::Tensor>& has_initial_state,
268
+ const torch::stable::Tensor& ssm_states, int64_t null_block_id,
269
+ int64_t block_size,
270
+ const std::optional<torch::stable::Tensor>& block_idx_first_scheduled_token,
271
+ const std::optional<torch::stable::Tensor>& block_idx_last_scheduled_token,
272
+ const std::optional<torch::stable::Tensor>& initial_state_idx,
273
+ const std::optional<torch::stable::Tensor>& cu_chunk_seqlen,
274
+ const std::optional<torch::stable::Tensor>& last_chunk_indices);
275
+
276
// Activation kernels (shared CUDA/ROCm)
277
void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
278
void silu_and_mul_clamp(torch::stable::Tensor& out,
csrc/libtorch_stable/sampler.cu
renamed
+59
-56
@@ -1,8 +1,6 @@
1
-#include "cuda_compat.h"
1
+#include "../cuda_compat.h"
2
#include "dispatch_utils.h"
3
-
4
-#include <torch/cuda.h>
5
-#include <c10/cuda/CUDAGuard.h>
3
+#include "torch_utils.h"
4
5
#ifndef USE_ROCM
6
#include <cub/cub.cuh>
@@ -618,14 +616,14 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(
616
} // namespace vllm
617
618
void apply_repetition_penalties_(
621
- torch::Tensor& logits, // [num_seqs, vocab_size], in-place
622
- const torch::Tensor& prompt_mask, // [num_seqs, vocab_size]
623
- const torch::Tensor& output_mask, // [num_seqs, vocab_size]
624
- const torch::Tensor& repetition_penalties) { // [num_seqs]
625
- TORCH_CHECK(logits.is_contiguous());
626
- TORCH_CHECK(prompt_mask.is_contiguous());
627
- TORCH_CHECK(output_mask.is_contiguous());
628
- TORCH_CHECK(repetition_penalties.is_contiguous());
619
+ torch::stable::Tensor& logits, // [num_seqs, vocab_size], in-place
620
+ const torch::stable::Tensor& prompt_mask, // [num_seqs, vocab_size]
621
+ const torch::stable::Tensor& output_mask, // [num_seqs, vocab_size]
622
+ const torch::stable::Tensor& repetition_penalties) { // [num_seqs]
623
+ STD_TORCH_CHECK(logits.is_contiguous());
624
+ STD_TORCH_CHECK(prompt_mask.is_contiguous());
625
+ STD_TORCH_CHECK(output_mask.is_contiguous());
626
+ STD_TORCH_CHECK(repetition_penalties.is_contiguous());
627
628
int vocab_size = logits.size(-1);
629
int num_seqs = logits.size(0);
@@ -635,7 +633,7 @@ void apply_repetition_penalties_(
633
// Get number of SMs on the current device
634
int sms = 0;
635
cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount,
638
- logits.get_device());
636
+ logits.get_device_index());
637
638
// Compute tile_num and tile_size
639
int tile_num =
@@ -645,27 +643,29 @@ void apply_repetition_penalties_(
643
// Each block handles one sequence and a tile of vocab
644
dim3 grid(num_seqs, tile_num);
645
dim3 block(std::min(tile_size, 1024));
648
- const at::cuda::OptionalCUDAGuard device_guard(device_of(logits));
649
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
650
- VLLM_DISPATCH_FLOATING_TYPES(
646
+ const torch::stable::accelerator::DeviceGuard device_guard(
647
+ logits.get_device_index());
648
+ const cudaStream_t stream = get_current_cuda_stream();
649
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
650
logits.scalar_type(), "apply_repetition_penalties_kernel", [&] {
651
vllm::apply_repetition_penalties_kernel<scalar_t>
652
<<<grid, block, 0, stream>>>(
654
- logits.data_ptr<scalar_t>(), prompt_mask.data_ptr<bool>(),
655
- output_mask.data_ptr<bool>(),
656
- repetition_penalties.data_ptr<scalar_t>(), num_seqs, vocab_size,
657
- tile_size);
653
+ logits.mutable_data_ptr<scalar_t>(),
654
+ prompt_mask.const_data_ptr<bool>(),
655
+ output_mask.const_data_ptr<bool>(),
656
+ repetition_penalties.const_data_ptr<scalar_t>(), num_seqs,
657
+ vocab_size, tile_size);
658
});
659
}
660
661
-void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,
662
- const torch::Tensor& seqLens, torch::Tensor& indices,
663
- int64_t numRows, int64_t stride0, int64_t stride1,
664
- int64_t topK) {
661
+void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n,
662
+ const torch::stable::Tensor& seqLens,
663
+ torch::stable::Tensor& indices, int64_t numRows,
664
+ int64_t stride0, int64_t stride1, int64_t topK) {
665
constexpr int kSortingAlgorithmThreshold = 12288;
666
constexpr int kSplitWorkThreshold = 200 * 1000;
667
constexpr int kNumThreadsPerBlock = 512;
668
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
668
+ const cudaStream_t stream = get_current_cuda_stream();
669
const auto numColumns = logits.size(1);
670
671
// True if seqLens is 2D (B, next_n): each logit row has its own pre-computed
@@ -677,73 +677,76 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,
677
// Use insertion sort
678
vllm::topKPerRowDecode<kNumThreadsPerBlock, false>
679
<<<numRows, kNumThreadsPerBlock, topK * sizeof(int32_t), stream>>>(
680
- logits.data_ptr<float>(), seqLens.data_ptr<int>(),
681
- indices.data_ptr<int>(), static_cast<int>(stride0),
680
+ logits.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
681
+ indices.mutable_data_ptr<int>(), static_cast<int>(stride0),
682
static_cast<int>(stride1), static_cast<int>(topK),
683
static_cast<int>(next_n), seqLensIs2D);
684
} else if (numColumns < kSplitWorkThreshold) {
685
// From this threshold, use radix sort instead
686
vllm::topKPerRowDecode<kNumThreadsPerBlock, true>
687
<<<numRows, kNumThreadsPerBlock, topK * sizeof(int32_t), stream>>>(
688
- logits.data_ptr<float>(), seqLens.data_ptr<int>(),
689
- indices.data_ptr<int>(), static_cast<int>(stride0),
688
+ logits.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
689
+ indices.mutable_data_ptr<int>(), static_cast<int>(stride0),
690
static_cast<int>(stride1), static_cast<int>(topK),
691
static_cast<int>(next_n), seqLensIs2D);
692
} else {
693
// Long sequences are run in two steps
694
constexpr auto multipleBlocksPerRowConfig = 10;
695
696
- const auto outIndicesAux =
697
- torch::empty({numRows, multipleBlocksPerRowConfig, topK},
698
- torch::dtype(torch::kInt32).device(logits.device()));
699
- const auto outLogitsAux =
700
- torch::empty({numRows, multipleBlocksPerRowConfig, topK},
701
- torch::dtype(torch::kFloat).device(logits.device()));
696
+ const auto outIndicesAux = torch::stable::empty(
697
+ {numRows, multipleBlocksPerRowConfig, topK},
698
+ torch::headeronly::ScalarType::Int, std::nullopt, logits.device());
699
+ const auto outLogitsAux = torch::stable::empty(
700
+ {numRows, multipleBlocksPerRowConfig, topK},
701
+ torch::headeronly::ScalarType::Float, std::nullopt, logits.device());
702
703
vllm::topKPerRowDecode<kNumThreadsPerBlock, true, true>
704
<<<dim3(numRows, multipleBlocksPerRowConfig), kNumThreadsPerBlock,
705
2 * topK * sizeof(int32_t), stream>>>(
706
- logits.data_ptr<float>(), seqLens.data_ptr<int>(),
707
- outIndicesAux.data_ptr<int>(), static_cast<int>(stride0),
706
+ logits.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
707
+ outIndicesAux.mutable_data_ptr<int>(), static_cast<int>(stride0),
708
static_cast<int>(stride1), static_cast<int>(topK),
709
static_cast<int>(next_n), seqLensIs2D,
710
- outLogitsAux.data_ptr<float>());
710
+ outLogitsAux.mutable_data_ptr<float>());
711
712
constexpr int kNumThreadsPerBlockMerge = 1024;
713
vllm::topKPerRowDecode<kNumThreadsPerBlockMerge, true, false, true>
714
<<<numRows, kNumThreadsPerBlockMerge, topK * sizeof(int32_t), stream>>>(
715
- outLogitsAux.data_ptr<float>(), seqLens.data_ptr<int>(),
716
- indices.data_ptr<int>(), multipleBlocksPerRowConfig * topK, 1,
717
- static_cast<int>(topK), static_cast<int>(next_n), seqLensIs2D,
718
- nullptr, multipleBlocksPerRowConfig, outIndicesAux.data_ptr<int>());
715
+ outLogitsAux.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
716
+ indices.mutable_data_ptr<int>(), multipleBlocksPerRowConfig * topK,
717
+ 1, static_cast<int>(topK), static_cast<int>(next_n), seqLensIs2D,
718
+ nullptr, multipleBlocksPerRowConfig,
719
+ outIndicesAux.const_data_ptr<int>());
720
}
721
}
722
722
-void top_k_per_row_prefill(const torch::Tensor& logits,
723
- const torch::Tensor& rowStarts,
724
- const torch::Tensor& rowEnds, torch::Tensor& indices,
725
- int64_t numRows, int64_t stride0, int64_t stride1,
726
- int64_t topK) {
723
+void top_k_per_row_prefill(const torch::stable::Tensor& logits,
724
+ const torch::stable::Tensor& rowStarts,
725
+ const torch::stable::Tensor& rowEnds,
726
+ torch::stable::Tensor& indices, int64_t numRows,
727
+ int64_t stride0, int64_t stride1, int64_t topK) {
728
constexpr int kSortingAlgorithmThreshold = 12288;
729
constexpr int kNumThreadsPerBlock = 512;
729
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
730
+ const cudaStream_t stream = get_current_cuda_stream();
731
732
int numInsertionBlocks =
733
std::min(static_cast<int>(numRows), kSortingAlgorithmThreshold);
734
vllm::topKPerRowPrefill<kNumThreadsPerBlock, false>
735
<<<numInsertionBlocks, kNumThreadsPerBlock, topK * sizeof(int32_t),
735
- stream>>>(logits.data_ptr<float>(), rowStarts.data_ptr<int>(),
736
- rowEnds.data_ptr<int>(), indices.data_ptr<int>(),
737
- static_cast<int>(stride0), static_cast<int>(stride1),
738
- static_cast<int>(topK), 0);
736
+ stream>>>(logits.const_data_ptr<float>(),
737
+ rowStarts.const_data_ptr<int>(),
738
+ rowEnds.const_data_ptr<int>(),
739
+ indices.mutable_data_ptr<int>(), static_cast<int>(stride0),
740
+ static_cast<int>(stride1), static_cast<int>(topK), 0);
741
742
if (numRows > kSortingAlgorithmThreshold) {
743
int numRadixBlocks = numRows - kSortingAlgorithmThreshold;
744
vllm::topKPerRowPrefill<kNumThreadsPerBlock, true>
745
<<<numRadixBlocks, kNumThreadsPerBlock, topK * sizeof(int32_t),
744
- stream>>>(logits.data_ptr<float>(), rowStarts.data_ptr<int>(),
745
- rowEnds.data_ptr<int>(), indices.data_ptr<int>(),
746
- static_cast<int>(stride0), static_cast<int>(stride1),
747
- static_cast<int>(topK), kSortingAlgorithmThreshold);
746
+ stream>>>(
747
+ logits.const_data_ptr<float>(), rowStarts.const_data_ptr<int>(),
748
+ rowEnds.const_data_ptr<int>(), indices.mutable_data_ptr<int>(),
749
+ static_cast<int>(stride0), static_cast<int>(stride1),
750
+ static_cast<int>(topK), kSortingAlgorithmThreshold);
751
}
752
}
\ No newline at end of file
csrc/libtorch_stable/topk.cu
renamed
+78
-69
@@ -1,49 +1,51 @@
1
// Persistent TopK kernel for DeepSeek V3 sparse attention indexer.
2
// See persistent_topk.cuh for kernel implementation.
3
4
-#include <torch/all.h>
5
-#include <ATen/cuda/CUDAContext.h>
4
#include <cuda_runtime.h>
5
#include <algorithm>
6
7
+#include "torch_utils.h"
8
+
9
#ifndef USE_ROCM
10
- #include "persistent_topk.cuh"
10
+ #include "../persistent_topk.cuh"
11
#endif
12
13
namespace {
14
15
#ifndef USE_ROCM
16
template <int TopK>
17
-void launch_persistent_topk(const torch::Tensor& logits,
18
- const torch::Tensor& lengths, torch::Tensor& output,
19
- torch::Tensor& workspace, int64_t max_seq_len) {
17
+void launch_persistent_topk(const torch::stable::Tensor& logits,
18
+ const torch::stable::Tensor& lengths,
19
+ torch::stable::Tensor& output,
20
+ torch::stable::Tensor& workspace,
21
+ int64_t max_seq_len) {
22
namespace P = vllm::persistent;
23
24
const int64_t num_rows = logits.size(0);
25
const int64_t stride = logits.stride(0);
24
- cudaStream_t stream = at::cuda::getCurrentCUDAStream();
26
+ const cudaStream_t stream = get_current_cuda_stream();
27
28
static int num_sms = 0;
29
static int max_smem_per_block = 0;
30
if (num_sms == 0) {
29
- int device;
30
- cudaGetDevice(&device);
31
- cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, device);
32
- cudaDeviceGetAttribute(&max_smem_per_block,
33
- cudaDevAttrMaxSharedMemoryPerBlockOptin, device);
31
+ const cudaDeviceProp* device_prop = get_device_prop();
32
+ num_sms = device_prop->multiProcessorCount;
33
+ max_smem_per_block = device_prop->sharedMemPerBlockOptin;
34
}
35
36
if (num_rows > 32 && max_smem_per_block >= 128 * 1024) {
37
cudaError_t status =
38
vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
39
- logits.data_ptr<float>(), output.data_ptr<int32_t>(),
40
- lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
39
+ logits.const_data_ptr<float>(), output.mutable_data_ptr<int32_t>(),
40
+ lengths.const_data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
41
static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride), stream);
42
- TORCH_CHECK(status == cudaSuccess,
43
- "FilteredTopK failed: ", cudaGetErrorString(status));
42
+ STD_TORCH_CHECK(status == cudaSuccess,
43
+ "FilteredTopK failed: ", cudaGetErrorString(status));
44
} else {
45
- TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
46
- TORCH_CHECK(workspace.dtype() == torch::kUInt8, "workspace must be uint8");
45
+ STD_TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
46
+ STD_TORCH_CHECK(
47
+ workspace.scalar_type() == torch::headeronly::ScalarType::Byte,
48
+ "workspace must be uint8");
49
50
int effective_max_smem;
51
if (num_rows <= 4) {
@@ -99,9 +101,9 @@ void launch_persistent_topk(const torch::Tensor& logits,
101
&occupancy, P::persistent_topk_kernel<TopK, 1>, P::kThreadsPerBlock,
102
smem_size);
103
}
102
- TORCH_CHECK(occ_err == cudaSuccess,
103
- "persistent_topk occupancy query failed: ",
104
- cudaGetErrorString(occ_err));
104
+ STD_TORCH_CHECK(occ_err == cudaSuccess,
105
+ "persistent_topk occupancy query failed: ",
106
+ cudaGetErrorString(occ_err));
107
if (occupancy < 1) occupancy = 1;
108
109
// The cooperative spin-wait barrier only runs when at least one row hits
@@ -131,27 +133,29 @@ void launch_persistent_topk(const torch::Tensor& logits,
133
// If the cooperative launch wouldn't fit, fall back to FilteredTopK
134
// instead of deadlocking. Only relevant when needs_cooperative.
135
if (needs_cooperative && total_ctas > hw_resident_cap) {
134
- TORCH_CHECK(max_smem_per_block >= 128 * 1024,
135
- "persistent_topk would oversubscribe and the FilteredTopK "
136
- "fallback requires >=128KB smem per block (have ",
137
- max_smem_per_block, "). total_ctas=", total_ctas,
138
- " > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
139
- ", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
140
- ", smem=", smem_size, ").");
136
+ STD_TORCH_CHECK(
137
+ max_smem_per_block >= 128 * 1024,
138
+ "persistent_topk would oversubscribe and the FilteredTopK "
139
+ "fallback requires >=128KB smem per block (have ",
140
+ max_smem_per_block, "). total_ctas=", total_ctas,
141
+ " > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
142
+ ", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
143
+ ", smem=", smem_size, ").");
144
cudaError_t status =
145
vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
143
- logits.data_ptr<float>(), output.data_ptr<int32_t>(),
144
- lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
145
- static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride),
146
- stream);
147
- TORCH_CHECK(status == cudaSuccess,
148
- "FilteredTopK fallback failed: ", cudaGetErrorString(status));
146
+ logits.const_data_ptr<float>(),
147
+ output.mutable_data_ptr<int32_t>(),
148
+ lengths.const_data_ptr<int32_t>(),
149
+ static_cast<uint32_t>(num_rows), static_cast<uint32_t>(TopK),
150
+ static_cast<uint32_t>(stride), stream);
151
+ STD_TORCH_CHECK(status == cudaSuccess, "FilteredTopK fallback failed: ",
152
+ cudaGetErrorString(status));
153
return;
154
}
155
156
size_t state_bytes = num_groups * sizeof(P::RadixRowState);
153
- TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
154
- "workspace too small, need ", state_bytes, " bytes");
157
+ STD_TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
158
+ "workspace too small, need ", state_bytes, " bytes");
159
160
// Zero the per-group RadixRowState region before launch.
161
//
@@ -179,22 +183,22 @@ void launch_persistent_topk(const torch::Tensor& logits,
183
// first red_release. cudaMemsetAsync is stream-ordered: the zero
184
// is globally visible before any CTA runs.
185
{
182
- cudaError_t mz_err = cudaMemsetAsync(workspace.data_ptr<uint8_t>(), 0,
183
- state_bytes, stream);
184
- TORCH_CHECK(mz_err == cudaSuccess,
185
- "row_states memset failed: ", cudaGetErrorString(mz_err));
186
+ cudaError_t mz_err = cudaMemsetAsync(
187
+ workspace.mutable_data_ptr<uint8_t>(), 0, state_bytes, stream);
188
+ STD_TORCH_CHECK(mz_err == cudaSuccess,
189
+ "row_states memset failed: ", cudaGetErrorString(mz_err));
190
}
191
192
P::PersistentTopKParams params;
189
- params.input = logits.data_ptr<float>();
190
- params.output = output.data_ptr<int32_t>();
191
- params.lengths = lengths.data_ptr<int32_t>();
193
+ params.input = logits.const_data_ptr<float>();
194
+ params.output = output.mutable_data_ptr<int32_t>();
195
+ params.lengths = lengths.const_data_ptr<int32_t>();
196
params.num_rows = static_cast<uint32_t>(num_rows);
197
params.stride = static_cast<uint32_t>(stride);
198
params.top_k = static_cast<uint32_t>(TopK);
199
params.chunk_size = chunk_size;
196
- params.row_states =
197
- reinterpret_cast<P::RadixRowState*>(workspace.data_ptr<uint8_t>());
200
+ params.row_states = reinterpret_cast<P::RadixRowState*>(
201
+ workspace.mutable_data_ptr<uint8_t>());
202
params.ctas_per_group = ctas_per_group;
203
params.max_seq_len = static_cast<uint32_t>(max_seq_len);
204
@@ -203,8 +207,8 @@ void launch_persistent_topk(const torch::Tensor& logits,
207
auto kernel = &P::persistent_topk_kernel<TOPK_VAL, VS>; \
208
cudaError_t err = cudaFuncSetAttribute( \
209
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); \
206
- TORCH_CHECK(err == cudaSuccess, \
207
- "Failed to set smem: ", cudaGetErrorString(err)); \
210
+ STD_TORCH_CHECK(err == cudaSuccess, \
211
+ "Failed to set smem: ", cudaGetErrorString(err)); \
212
kernel<<<total_ctas, P::kThreadsPerBlock, smem_size, stream>>>(params); \
213
} while (0)
214
@@ -219,37 +223,42 @@ void launch_persistent_topk(const torch::Tensor& logits,
223
}
224
225
cudaError_t err = cudaGetLastError();
222
- TORCH_CHECK(err == cudaSuccess,
223
- "persistent_topk failed: ", cudaGetErrorString(err));
226
+ STD_TORCH_CHECK(err == cudaSuccess,
227
+ "persistent_topk failed: ", cudaGetErrorString(err));
228
}
229
#endif
230
231
} // anonymous namespace
232
229
-void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
230
- torch::Tensor& output, torch::Tensor& workspace, int64_t k,
233
+void persistent_topk(const torch::stable::Tensor& logits,
234
+ const torch::stable::Tensor& lengths,
235
+ torch::stable::Tensor& output,
236
+ torch::stable::Tensor& workspace, int64_t k,
237
int64_t max_seq_len) {
238
#ifndef USE_ROCM
233
- TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
234
- TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
235
- TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
236
- TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
237
- TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
238
- TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
239
- TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
240
- TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
241
- "lengths must be 1D or 2D");
242
- TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
243
- TORCH_CHECK(output.dim() == 2, "output must be 2D");
239
+ STD_TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
240
+ STD_TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
241
+ STD_TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
242
+ STD_TORCH_CHECK(logits.scalar_type() == torch::headeronly::ScalarType::Float,
243
+ "Only float32 supported");
244
+ STD_TORCH_CHECK(lengths.scalar_type() == torch::headeronly::ScalarType::Int,
245
+ "lengths must be int32");
246
+ STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Int,
247
+ "output must be int32");
248
+ STD_TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
249
+ STD_TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
250
+ "lengths must be 1D or 2D");
251
+ STD_TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
252
+ STD_TORCH_CHECK(output.dim() == 2, "output must be 2D");
253
254
const int64_t num_rows = logits.size(0);
246
- const int64_t stride = logits.stride(0);
255
248
- TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
249
- TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
250
- "output size mismatch");
251
- TORCH_CHECK(k == 512 || k == 1024 || k == 2048,
252
- "persistent_topk supports k=512, k=1024, or k=2048, got k=", k);
256
+ STD_TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
257
+ STD_TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
258
+ "output size mismatch");
259
+ STD_TORCH_CHECK(
260
+ k == 512 || k == 1024 || k == 2048,
261
+ "persistent_topk supports k=512, k=1024, or k=2048, got k=", k);
262
263
if (k == 512) {
264
launch_persistent_topk<512>(logits, lengths, output, workspace,
@@ -262,6 +271,6 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
271
max_seq_len);
272
}
273
#else
265
- TORCH_CHECK(false, "persistent_topk is not supported on ROCm");
274
+ STD_TORCH_CHECK(false, "persistent_topk is not supported on ROCm");
275
#endif
276
}
csrc/libtorch_stable/torch_bindings.cpp
+62
@@ -263,6 +263,20 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
263
"CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor");
264
#endif
265
266
+ // Merge attn states
267
+ // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
268
+ // can be used to combine partial attention results (in the split-KV case)
269
+ ops.def(
270
+ "merge_attn_states("
271
+ " Tensor! output,"
272
+ " Tensor!? output_lse,"
273
+ " Tensor prefix_output,"
274
+ " Tensor prefix_lse,"
275
+ " Tensor suffix_output,"
276
+ " Tensor suffix_lse,"
277
+ " int!? prefill_tokens_with_context,"
278
+ " Tensor? output_scale=None) -> ()");
279
+
280
// Hadamard transforms
281
// conditionally compiled so impl registration is in source file
282
ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor");
@@ -319,6 +333,26 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
333
"bool is_neox, Tensor position_ids, "
334
"int forced_token_heads_per_warp=-1) -> ()");
335
336
+ // Apply repetition penalties to logits in-place.
337
+ ops.def(
338
+ "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
339
+ "Tensor output_mask, Tensor repetition_penalties) -> ()");
340
+
341
+ // Optimized top-k per row operations.
342
+ ops.def(
343
+ "top_k_per_row_prefill(Tensor logits, Tensor rowStarts, Tensor rowEnds, "
344
+ "Tensor! indices, int numRows, int stride0, "
345
+ "int stride1, int topK) -> ()");
346
+
347
+ ops.def(
348
+ "top_k_per_row_decode(Tensor logits, int next_n, "
349
+ "Tensor seq_lens, Tensor! indices, "
350
+ "int numRows, int stride0, int stride1, int topK) -> ()");
351
+
352
+ ops.def(
353
+ "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, "
354
+ "Tensor workspace, int k, int max_seq_len) -> ()");
355
+
356
// Activation ops
357
// Activation function used in SwiGLU.
358
ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
@@ -422,6 +456,24 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
456
"int type, SymInt row, SymInt tokens) -> Tensor");
457
458
ops.def("ggml_moe_get_block_size(int type) -> int");
459
+
460
+ // Mamba selective scan kernel
461
+ ops.def(
462
+ "selective_scan_fwd(Tensor! u, Tensor! delta,"
463
+ "Tensor! A, Tensor! B, Tensor! C,"
464
+ "Tensor? D_, Tensor!? z_, Tensor? delta_bias_,"
465
+ "bool delta_softplus,"
466
+ "Tensor? query_start_loc,"
467
+ "Tensor? cache_indices,"
468
+ "Tensor? has_initial_state,"
469
+ "Tensor! ssm_states,"
470
+ "int null_block_id,"
471
+ "int block_size,"
472
+ "Tensor? block_idx_first_scheduled_token,"
473
+ "Tensor? block_idx_last_scheduled_token,"
474
+ "Tensor? initial_state_idx,"
475
+ "Tensor? cu_chunk_seqlen,"
476
+ "Tensor? last_chunk_indices) -> ()");
477
}
478
479
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
@@ -469,6 +521,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
521
// files (allspark_repack.cu and allspark_qgemm_w8a16.cu)
522
#endif
523
524
+ ops.impl("merge_attn_states", TORCH_BOX(&merge_attn_states));
525
+
526
// Layernorm kernels (shared CUDA/ROCm)
527
ops.impl("rms_norm", TORCH_BOX(&rms_norm));
528
ops.impl("fused_add_rms_norm", TORCH_BOX(&fused_add_rms_norm));
@@ -487,6 +541,13 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
541
ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding));
542
ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope));
543
544
+ // Sampler kernels (shared CUDA/ROCm)
545
+ ops.impl("apply_repetition_penalties_",
546
+ TORCH_BOX(&apply_repetition_penalties_));
547
+ ops.impl("top_k_per_row_prefill", TORCH_BOX(&top_k_per_row_prefill));
548
+ ops.impl("top_k_per_row_decode", TORCH_BOX(&top_k_per_row_decode));
549
+ ops.impl("persistent_topk", TORCH_BOX(&persistent_topk));
550
+
551
// Activation kernels (shared CUDA/ROCm)
552
ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul));
553
ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu));
@@ -519,6 +580,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
580
ops.impl("ggml_mul_mat_a8", TORCH_BOX(&ggml_mul_mat_a8));
581
ops.impl("ggml_moe_a8", TORCH_BOX(&ggml_moe_a8));
582
ops.impl("ggml_moe_a8_vec", TORCH_BOX(&ggml_moe_a8_vec));
583
+ ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd));
584
}
585
586
// These capability-check functions take only primitive args (no tensors), so
csrc/ops.h
-43
@@ -54,13 +54,6 @@ void paged_attention_v2(
54
const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
55
const int64_t blocksparse_head_sliding_step);
56
57
-void merge_attn_states(
58
- torch::Tensor& output, std::optional<torch::Tensor> output_lse,
59
- const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse,
60
- const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse,
61
- const std::optional<int64_t> prefill_tokens_with_context,
62
- const std::optional<torch::Tensor>& output_scale = std::nullopt);
63
-
57
// rms_norm and fused_add_rms_norm declarations also exist in
58
// csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here
59
// because the CPU build still uses these torch::Tensor declarations.
@@ -76,26 +69,6 @@ torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
69
torch::Tensor const& cos_sin_cache, int64_t q_head_padded, double eps,
70
int64_t cache_block_size);
71
79
-void apply_repetition_penalties_(torch::Tensor& logits,
80
- const torch::Tensor& prompt_mask,
81
- const torch::Tensor& output_mask,
82
- const torch::Tensor& repetition_penalties);
83
-
84
-void top_k_per_row_prefill(const torch::Tensor& logits,
85
- const torch::Tensor& rowStarts,
86
- const torch::Tensor& rowEnds, torch::Tensor& indices,
87
- int64_t numRows, int64_t stride0, int64_t stride1,
88
- int64_t topK);
89
-
90
-void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,
91
- const torch::Tensor& seqLens, torch::Tensor& indices,
92
- int64_t numRows, int64_t stride0, int64_t stride1,
93
- int64_t topK);
94
-
95
-void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
96
- torch::Tensor& output, torch::Tensor& workspace, int64_t k,
97
- int64_t max_seq_len);
98
-
72
void silu_and_mul_per_block_quant(torch::Tensor& out,
73
torch::Tensor const& input,
74
torch::Tensor& scales, int64_t group_size,
@@ -150,22 +123,6 @@ void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
123
torch::Tensor& scales,
124
std::optional<torch::Tensor> const& azp);
125
153
-void selective_scan_fwd(
154
- const torch::Tensor& u, const torch::Tensor& delta, const torch::Tensor& A,
155
- const torch::Tensor& B, const torch::Tensor& C,
156
- const std::optional<torch::Tensor>& D_,
157
- const std::optional<torch::Tensor>& z_,
158
- const std::optional<torch::Tensor>& delta_bias_, bool delta_softplus,
159
- const std::optional<torch::Tensor>& query_start_loc,
160
- const std::optional<torch::Tensor>& cache_indices,
161
- const std::optional<torch::Tensor>& has_initial_state,
162
- const torch::Tensor& ssm_states, int64_t null_block_id, int64_t block_size,
163
- const std::optional<torch::Tensor>& block_idx_first_scheduled_token,
164
- const std::optional<torch::Tensor>& block_idx_last_scheduled_token,
165
- const std::optional<torch::Tensor>& initial_state_idx,
166
- const std::optional<torch::Tensor>& cu_chunk_seqlen,
167
- const std::optional<torch::Tensor>& last_chunk_indices);
168
-
126
torch::Tensor dynamic_4bit_int_moe_cpu(
127
torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights,
128
torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I,
csrc/persistent_topk.cuh
+9
-7
@@ -126,10 +126,10 @@ struct RadixRowState {
126
// ============================================================================
127
128
struct PersistentTopKParams {
129
- const float* __restrict__ input; // [num_rows, stride]
130
- int32_t* __restrict__ output; // [num_rows, top_k]
131
- int32_t* __restrict__ lengths; // [num_rows]
132
- RadixRowState* row_states; // large path: per-group state
129
+ const float* __restrict__ input; // [num_rows, stride]
130
+ int32_t* __restrict__ output; // [num_rows, top_k]
131
+ const int32_t* __restrict__ lengths; // [num_rows]
132
+ RadixRowState* row_states; // large path: per-group state
133
uint32_t num_rows;
134
uint32_t stride;
135
uint32_t top_k; // actual k value for output stride
@@ -1269,9 +1269,11 @@ constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) {
1269
}
1270
1271
template <typename DType, typename IdType, uint32_t MAX_K = 2048>
1272
-cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,
1273
- IdType* lengths, uint32_t num_rows,
1274
- uint32_t top_k_val, uint32_t max_len,
1272
+cudaError_t FilteredTopKRaggedTransform(const DType* input,
1273
+ IdType* output_indices,
1274
+ const IdType* lengths,
1275
+ uint32_t num_rows, uint32_t top_k_val,
1276
+ uint32_t max_len,
1277
cudaStream_t stream = 0) {
1278
constexpr size_t smem_size = FILTERED_TOPK_SMEM_DYNAMIC;
1279
constexpr int MAX_VEC = 16 / sizeof(DType);
csrc/torch_bindings.cpp
-59
@@ -62,21 +62,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
62
" int blocksparse_head_sliding_step) -> ()");
63
ops.impl("paged_attention_v2", torch::kCUDA, &paged_attention_v2);
64
65
- // Merge attn states
66
- // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
67
- // can be used to combine partial attention results (in the split-KV case)
68
- ops.def(
69
- "merge_attn_states("
70
- " Tensor! output,"
71
- " Tensor!? output_lse,"
72
- " Tensor prefix_output,"
73
- " Tensor prefix_lse,"
74
- " Tensor suffix_output,"
75
- " Tensor suffix_lse,"
76
- " int!? prefill_tokens_with_context,"
77
- " Tensor? output_scale=None) -> ()");
78
- ops.impl("merge_attn_states", torch::kCUDA, &merge_attn_states);
79
-
65
// Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
66
ops.def(
67
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
@@ -105,31 +90,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
90
ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
91
&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);
92
108
- // Apply repetition penalties to logits in-place
109
- ops.def(
110
- "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
111
- "Tensor output_mask, Tensor repetition_penalties) -> ()");
112
- ops.impl("apply_repetition_penalties_", torch::kCUDA,
113
- &apply_repetition_penalties_);
114
-
115
- // Optimized top-k per row operation
116
- ops.def(
117
- "top_k_per_row_prefill(Tensor logits, Tensor rowStarts, Tensor rowEnds, "
118
- "Tensor! indices, int numRows, int stride0, "
119
- "int stride1, int topK) -> ()");
120
- ops.impl("top_k_per_row_prefill", torch::kCUDA, &top_k_per_row_prefill);
121
-
122
- ops.def(
123
- "top_k_per_row_decode(Tensor logits, int next_n, "
124
- "Tensor seq_lens, Tensor! indices, "
125
- "int numRows, int stride0, int stride1, int topK) -> ()");
126
- ops.impl("top_k_per_row_decode", torch::kCUDA, &top_k_per_row_decode);
127
-
128
- ops.def(
129
- "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, "
130
- "Tensor workspace, int k, int max_seq_len) -> ()");
131
- ops.impl("persistent_topk", torch::kCUDA, &persistent_topk);
132
-
93
// Quantization ops
94
#ifndef USE_ROCM
95
@@ -230,25 +190,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
190
191
#endif
192
233
- // Mamba selective scan kernel
234
- ops.def(
235
- "selective_scan_fwd(Tensor! u, Tensor! delta,"
236
- "Tensor! A, Tensor! B, Tensor! C,"
237
- "Tensor? D_, Tensor!? z_, Tensor? delta_bias_,"
238
- "bool delta_softplus,"
239
- "Tensor? query_start_loc,"
240
- "Tensor? cache_indices,"
241
- "Tensor? has_initial_state,"
242
- "Tensor! ssm_states,"
243
- "int null_block_id,"
244
- "int block_size,"
245
- "Tensor? block_idx_first_scheduled_token,"
246
- "Tensor? block_idx_last_scheduled_token,"
247
- "Tensor? initial_state_idx,"
248
- "Tensor? cu_chunk_seqlen,"
249
- "Tensor? last_chunk_indices) -> ()");
250
- ops.impl("selective_scan_fwd", torch::kCUDA, &selective_scan_fwd);
251
-
193
#ifndef USE_ROCM
194
ops.def(
195
"minimax_allreduce_rms("
pyproject.toml
+1
@@ -141,6 +141,7 @@ bbc5b7ede = "bbc5b7ede"
141
NOOPs = "NOOPs"
142
nin_shortcut = "nin_shortcut"
143
cudaDevAttrMaxSharedMemoryPerBlockOptin = "cudaDevAttrMaxSharedMemoryPerBlockOptin"
144
+sharedMemPerBlockOptin = "sharedMemPerBlockOptin"
145
146
depthwise_seperable_out_channel = "depthwise_seperable_out_channel"
147
pard_token = "pard_token"