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, 06:00 UTC
[CPU][Perf] Enable fused kernels for GDN's gated delta rules (#43534)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Fadi Arafeh committed
Jun 2, 2026 at 09:00 UTC
0b25cf441903b06811e0f75f09c303c18fe08240
11 files changed
+812
-585
.buildkite/hardware_tests/cpu.yaml
+3
-1
@@ -16,6 +16,7 @@ steps:
16
- tests/kernels/test_onednn.py
17
- tests/kernels/test_awq_int4_to_int8.py
18
- tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
19
+ - tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
20
commands:
21
- |
22
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
@@ -24,7 +25,8 @@ steps:
25
pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py
26
pytest -x -v -s tests/kernels/test_onednn.py
27
pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
27
- pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py"
28
+ pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
29
+ pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
30
31
- label: CPU-Compatibility Tests
32
depends_on: []
.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
+2
-1
@@ -37,7 +37,8 @@ function cpu_tests() {
37
pytest -x -v -s tests/kernels/test_onednn.py
38
pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
39
pytest -x -v -s tests/kernels/core/test_cpu_activation.py
40
- pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic"
40
+ pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic
41
+ pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
42
43
# skip tests requiring model downloads if HF_TOKEN is not set
44
# due to rate-limits
cmake/cpu_extension.cmake
+18
-1
@@ -369,6 +369,18 @@ else()
369
add_compile_definitions(-DVLLM_NUMA_DISABLED)
370
endif()
371
372
+# check if the pytorch wheel ships libopenblas.so.
373
+set(VLLM_OPENBLAS_LIB "")
374
+if (NOT ENABLE_X86_ISA)
375
+ file(GLOB _VLLM_TORCH_OPENBLAS_LIBS
376
+ "${TORCH_INSTALL_PREFIX}/lib/libopenblas*.so*")
377
+ # Note: we don't link openblas directly to _C extension, as it's available through libtorch.so
378
+ if (_VLLM_TORCH_OPENBLAS_LIBS)
379
+ list(GET _VLLM_TORCH_OPENBLAS_LIBS 0 VLLM_OPENBLAS_LIB)
380
+ message(STATUS "CPU OpenBLAS library: ${VLLM_OPENBLAS_LIB}")
381
+ endif()
382
+endif()
383
+
384
#
385
# Generate CPU attention dispatch header
386
#
@@ -387,6 +399,7 @@ endif()
399
#
400
set(VLLM_EXT_SRC
401
"csrc/cpu/activation.cpp"
402
+ "csrc/cpu/sgl-kernels/fla.cpp"
403
"csrc/cpu/utils.cpp"
404
"csrc/cpu/spec_decode_utils.cpp"
405
"csrc/cpu/layernorm.cpp"
@@ -418,7 +431,6 @@ endif()
431
432
if (ENABLE_X86_ISA)
433
set(VLLM_EXT_SRC_SGL
421
- "csrc/cpu/sgl-kernels/fla.cpp"
434
"csrc/cpu/sgl-kernels/conv.cpp"
435
"csrc/cpu/sgl-kernels/gemm.cpp"
436
"csrc/cpu/sgl-kernels/gemm_int8.cpp"
@@ -430,6 +442,7 @@ if (ENABLE_X86_ISA)
442
"csrc/cpu/sgl-kernels/moe_fp8.cpp")
443
444
set(VLLM_EXT_SRC_AVX512
445
+ "csrc/cpu/sgl-kernels/fla.cpp"
446
"csrc/cpu/shm.cpp"
447
"csrc/cpu/cpu_wna16.cpp"
448
"csrc/cpu/cpu_fused_moe.cpp"
@@ -446,6 +459,7 @@ if (ENABLE_X86_ISA)
459
"csrc/moe/dynamic_4bit_int_moe_cpu.cpp")
460
461
set(VLLM_EXT_SRC_AVX2
462
+ "csrc/cpu/sgl-kernels/fla.cpp"
463
"csrc/cpu/utils.cpp"
464
"csrc/cpu/spec_decode_utils.cpp"
465
"csrc/cpu/cpu_attn.cpp"
@@ -519,6 +533,9 @@ else()
533
USE_SABI 3
534
WITH_SOABI
535
)
536
+ if (VLLM_OPENBLAS_LIB)
537
+ target_compile_definitions(_C PRIVATE VLLM_HAS_OPENBLAS)
538
+ endif()
539
endif()
540
541
message(STATUS "Enabling C extension.")
csrc/cpu/sgl-kernels/blas_gemm.h
new
+82
@@ -0,0 +1,82 @@
1
+// SPDX-License-Identifier: Apache-2.0
2
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+#include <ATen/native/CPUBlas.h>
5
+
6
+// Unlike brgemm, PyTorch does not publicly expose at::native::cpublas::gemm
7
+// If OpenBLS is available in the PyTorch wheel, we rely on it for fast
8
+// bf16:bf16->fp32 GEMMs Otherwise, we fall back to PyTorch reference BLAS path.
9
+#if defined(VLLM_HAS_OPENBLAS)
10
+extern "C" void sbgemm_(char* transa, char* transb, int* m, int* n, int* k,
11
+ float* alpha, const at::BFloat16* a, int* lda,
12
+ const at::BFloat16* b, int* ldb, float* beta, float* c,
13
+ int* ldc);
14
+
15
+extern "C" void sgemm_(char* transa, char* transb, int* m, int* n, int* k,
16
+ float* alpha, const float* a, int* lda, const float* b,
17
+ int* ldb, float* beta, float* c, int* ldc);
18
+
19
+inline char blas_transpose(at::native::TransposeType trans) {
20
+ switch (trans) {
21
+ case at::native::TransposeType::NoTranspose:
22
+ return 'n';
23
+ case at::native::TransposeType::Transpose:
24
+ return 't';
25
+ case at::native::TransposeType::ConjTranspose:
26
+ return 'c';
27
+ }
28
+ return 'n';
29
+}
30
+
31
+inline void blas_gemm(at::native::TransposeType transa,
32
+ at::native::TransposeType transb, int64_t m, int64_t n,
33
+ int64_t k, float alpha, const at::BFloat16* a,
34
+ int64_t lda, const at::BFloat16* b, int64_t ldb,
35
+ float beta, float* c, int64_t ldc) {
36
+ char transa_ = blas_transpose(transa);
37
+ char transb_ = blas_transpose(transb);
38
+ int m_ = static_cast<int>(m);
39
+ int n_ = static_cast<int>(n);
40
+ int k_ = static_cast<int>(k);
41
+ int lda_ = static_cast<int>(lda);
42
+ int ldb_ = static_cast<int>(ldb);
43
+ int ldc_ = static_cast<int>(ldc);
44
+ sbgemm_(&transa_, &transb_, &m_, &n_, &k_, &alpha, a, &lda_, b, &ldb_, &beta,
45
+ c, &ldc_);
46
+}
47
+
48
+inline void blas_gemm(at::native::TransposeType transa,
49
+ at::native::TransposeType transb, int64_t m, int64_t n,
50
+ int64_t k, float alpha, const float* a, int64_t lda,
51
+ const float* b, int64_t ldb, float beta, float* c,
52
+ int64_t ldc) {
53
+ char transa_ = blas_transpose(transa);
54
+ char transb_ = blas_transpose(transb);
55
+ int m_ = static_cast<int>(m);
56
+ int n_ = static_cast<int>(n);
57
+ int k_ = static_cast<int>(k);
58
+ int lda_ = static_cast<int>(lda);
59
+ int ldb_ = static_cast<int>(ldb);
60
+ int ldc_ = static_cast<int>(ldc);
61
+ sgemm_(&transa_, &transb_, &m_, &n_, &k_, &alpha, a, &lda_, b, &ldb_, &beta,
62
+ c, &ldc_);
63
+}
64
+
65
+inline void blas_gemm(at::native::TransposeType, at::native::TransposeType,
66
+ int64_t, int64_t, int64_t, float, const at::Half*,
67
+ int64_t, const at::Half*, int64_t, float, float*,
68
+ int64_t) {
69
+ TORCH_CHECK(false, "CPU OpenBLAS hgemm is not available.");
70
+}
71
+#else
72
+template <typename scalar_t>
73
+inline void blas_gemm(at::native::TransposeType transa,
74
+ at::native::TransposeType transb, int64_t m, int64_t n,
75
+ int64_t k, float alpha, const scalar_t* a, int64_t lda,
76
+ const scalar_t* b, int64_t ldb, float beta, float* c,
77
+ int64_t ldc) {
78
+ auto gemm = at::native::cpublas::gemm_no_downcast_stub.DEFAULT;
79
+ gemm(c10::CppTypeToScalarType<scalar_t>::value, transa, transb, m, n, k,
80
+ at::Scalar(alpha), a, lda, b, ldb, at::Scalar(beta), c, ldc);
81
+}
82
+#endif
\ No newline at end of file
csrc/cpu/sgl-kernels/fla.cpp
+278
-141
@@ -301,25 +301,42 @@ void chunk_gated_delta_rule_kernel_impl(
301
// attn = k_beta @ key.transpose(-1, -2)
302
// attn: [B, HV, num_chunk, chunk_size, chunk_size]
303
// transpose and pack for key
304
- pack_vnni<scalar_t>(
305
- /* dst */ k_transpose,
306
- /* src */ curr_k_pad,
307
- /* N */ chunk_size,
308
- /* K */ qk_head_size,
309
- /* ld_src */ qk_head_size,
310
- /* ld_dst */ chunk_size);
311
- // k_beta @ key.transpose(-1, -2)
312
- at::native::cpublas::brgemm(
313
- /* M */ chunk_size,
314
- /* N */ chunk_size,
315
- /* K */ qk_head_size,
316
- /* lda */ qk_head_size,
317
- /* ldb */ chunk_size,
318
- /* ldc */ chunk_size,
319
- /* add_C */ false,
320
- /* A */ curr_k_beta,
321
- /* B */ k_transpose,
322
- /* C */ curr_attn);
304
+ if constexpr (brgemm_supported()) {
305
+ pack_vnni<scalar_t>(
306
+ /* dst */ k_transpose,
307
+ /* src */ curr_k_pad,
308
+ /* N */ chunk_size,
309
+ /* K */ qk_head_size,
310
+ /* ld_src */ qk_head_size,
311
+ /* ld_dst */ chunk_size);
312
+ // k_beta @ key.transpose(-1, -2)
313
+ at::native::cpublas::brgemm(
314
+ /* M */ chunk_size,
315
+ /* N */ chunk_size,
316
+ /* K */ qk_head_size,
317
+ /* lda */ qk_head_size,
318
+ /* ldb */ chunk_size,
319
+ /* ldc */ chunk_size,
320
+ /* add_C */ false,
321
+ /* A */ curr_k_beta,
322
+ /* B */ k_transpose,
323
+ /* C */ curr_attn);
324
+ } else {
325
+ blas_gemm(
326
+ at::native::TransposeType::Transpose,
327
+ at::native::TransposeType::NoTranspose,
328
+ chunk_size,
329
+ chunk_size,
330
+ qk_head_size,
331
+ 1.0f,
332
+ curr_k_pad,
333
+ qk_head_size,
334
+ curr_k_beta,
335
+ qk_head_size,
336
+ 0.0f,
337
+ curr_attn,
338
+ chunk_size);
339
+ }
340
// attn = attn * decay_mask
341
for (int64_t m = 0; m < chunk_size; m++) {
342
at::vec::map2<float>(
@@ -413,25 +430,42 @@ void chunk_gated_delta_rule_kernel_impl(
430
// k_beta_g = k_beta * g: [B, HV, num_chunk, chunk_size, EK]
431
// k_cumdecay: [B, HV, num_chunk, chunk_size, EK]
432
// pack for value
416
- pack_vnni2<scalar_t>(
417
- /* dst */ v_pack,
418
- /* src */ curr_v_beta,
419
- /* N */ chunk_size,
420
- /* K */ v_head_size,
421
- /* ld_src */ v_head_size,
422
- /* ld_dst */ v_head_size);
423
- // value = attn @ v_beta
424
- at::native::cpublas::brgemm(
425
- /* M */ chunk_size,
426
- /* N */ v_head_size,
427
- /* K */ chunk_size,
428
- /* lda */ chunk_size,
429
- /* ldb */ v_head_size,
430
- /* ldc */ v_head_size,
431
- /* add_C */ false,
432
- /* A */ curr_attn_reduced,
433
- /* B */ v_pack,
434
- /* C */ curr_value);
433
+ if constexpr (brgemm_supported()) {
434
+ pack_vnni2<scalar_t>(
435
+ /* dst */ v_pack,
436
+ /* src */ curr_v_beta,
437
+ /* N */ chunk_size,
438
+ /* K */ v_head_size,
439
+ /* ld_src */ v_head_size,
440
+ /* ld_dst */ v_head_size);
441
+ // value = attn @ v_beta
442
+ at::native::cpublas::brgemm(
443
+ /* M */ chunk_size,
444
+ /* N */ v_head_size,
445
+ /* K */ chunk_size,
446
+ /* lda */ chunk_size,
447
+ /* ldb */ v_head_size,
448
+ /* ldc */ v_head_size,
449
+ /* add_C */ false,
450
+ /* A */ curr_attn_reduced,
451
+ /* B */ v_pack,
452
+ /* C */ curr_value);
453
+ } else {
454
+ blas_gemm(
455
+ at::native::TransposeType::NoTranspose,
456
+ at::native::TransposeType::NoTranspose,
457
+ v_head_size,
458
+ chunk_size,
459
+ chunk_size,
460
+ 1.0f,
461
+ curr_v_beta,
462
+ v_head_size,
463
+ curr_attn_reduced,
464
+ chunk_size,
465
+ 0.0f,
466
+ curr_value,
467
+ v_head_size);
468
+ }
469
// k_beta_g = k_beta * g.exp().unsqueeze(-1)
470
for (int64_t j = 0; j < chunk_size; j++) {
471
int64_t i = 0;
@@ -445,25 +479,42 @@ void chunk_gated_delta_rule_kernel_impl(
479
}
480
}
481
// pack for k_beta_g
448
- pack_vnni2<scalar_t>(
449
- /* dst */ k_beta_g_pack,
450
- /* src */ k_beta_g,
451
- /* N */ chunk_size,
452
- /* K */ qk_head_size,
453
- /* ld_src */ qk_head_size,
454
- /* ld_dst */ qk_head_size);
455
- // k_cumdecay = attn @ k_beta_g
456
- at::native::cpublas::brgemm(
457
- /* M */ chunk_size,
458
- /* N */ qk_head_size,
459
- /* K */ chunk_size,
460
- /* lda */ chunk_size,
461
- /* ldb */ qk_head_size,
462
- /* ldc */ qk_head_size,
463
- /* add_C */ false,
464
- /* A */ curr_attn_reduced,
465
- /* B */ k_beta_g_pack,
466
- /* C */ k_cumdecay);
482
+ if constexpr (brgemm_supported()) {
483
+ pack_vnni2<scalar_t>(
484
+ /* dst */ k_beta_g_pack,
485
+ /* src */ k_beta_g,
486
+ /* N */ chunk_size,
487
+ /* K */ qk_head_size,
488
+ /* ld_src */ qk_head_size,
489
+ /* ld_dst */ qk_head_size);
490
+ // k_cumdecay = attn @ k_beta_g
491
+ at::native::cpublas::brgemm(
492
+ /* M */ chunk_size,
493
+ /* N */ qk_head_size,
494
+ /* K */ chunk_size,
495
+ /* lda */ chunk_size,
496
+ /* ldb */ qk_head_size,
497
+ /* ldc */ qk_head_size,
498
+ /* add_C */ false,
499
+ /* A */ curr_attn_reduced,
500
+ /* B */ k_beta_g_pack,
501
+ /* C */ k_cumdecay);
502
+ } else {
503
+ blas_gemm(
504
+ at::native::TransposeType::NoTranspose,
505
+ at::native::TransposeType::NoTranspose,
506
+ qk_head_size,
507
+ chunk_size,
508
+ chunk_size,
509
+ 1.0f,
510
+ k_beta_g,
511
+ qk_head_size,
512
+ curr_attn_reduced,
513
+ chunk_size,
514
+ 0.0f,
515
+ k_cumdecay,
516
+ qk_head_size);
517
+ }
518
for (int i = 0; i < chunk_size; i++) {
519
at::vec::map<scalar_t>(
520
[](fVec x) { return x; },
@@ -551,25 +602,42 @@ void chunk_gated_delta_rule_kernel_impl(
602
603
// attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0)
604
// k_transpose_i = k_i.transpose(-1, -2)
554
- pack_vnni<scalar_t>(
555
- /* dst */ k_transpose_i,
556
- /* src */ k_i,
557
- /* N */ chunk_size,
558
- /* K */ qk_head_size,
559
- /* ld_src */ qk_head_size,
560
- /* ld_dst */ chunk_size);
561
- // attn_i = q_i @ k_transpose_i
562
- at::native::cpublas::brgemm(
563
- /* M */ chunk_size,
564
- /* N */ chunk_size,
565
- /* K */ qk_head_size,
566
- /* lda */ qk_head_size,
567
- /* ldb */ chunk_size,
568
- /* ldc */ chunk_size,
569
- /* add_C */ false,
570
- /* A */ q_i,
571
- /* B */ k_transpose_i,
572
- /* C */ attn_i);
605
+ if constexpr (brgemm_supported()) {
606
+ pack_vnni<scalar_t>(
607
+ /* dst */ k_transpose_i,
608
+ /* src */ k_i,
609
+ /* N */ chunk_size,
610
+ /* K */ qk_head_size,
611
+ /* ld_src */ qk_head_size,
612
+ /* ld_dst */ chunk_size);
613
+ // attn_i = q_i @ k_transpose_i
614
+ at::native::cpublas::brgemm(
615
+ /* M */ chunk_size,
616
+ /* N */ chunk_size,
617
+ /* K */ qk_head_size,
618
+ /* lda */ qk_head_size,
619
+ /* ldb */ chunk_size,
620
+ /* ldc */ chunk_size,
621
+ /* add_C */ false,
622
+ /* A */ q_i,
623
+ /* B */ k_transpose_i,
624
+ /* C */ attn_i);
625
+ } else {
626
+ blas_gemm(
627
+ at::native::TransposeType::Transpose,
628
+ at::native::TransposeType::NoTranspose,
629
+ chunk_size,
630
+ chunk_size,
631
+ qk_head_size,
632
+ 1.0f,
633
+ k_i,
634
+ qk_head_size,
635
+ q_i,
636
+ qk_head_size,
637
+ 0.0f,
638
+ attn_i,
639
+ chunk_size);
640
+ }
641
// attn_i = attn_i * decay_mask_i
642
for (int64_t m = 0; m < chunk_size; m++) {
643
auto attn_i_m = attn_i + m * chunk_size;
@@ -609,28 +677,45 @@ void chunk_gated_delta_rule_kernel_impl(
677
}
678
679
// pack for curr_last_recurrent_state
612
- pack_vnni2<scalar_t>(
613
- /* dst */ curr_last_recurrent_state_pack_reduced,
614
- /* src */ curr_last_recurrent_state_reduced,
615
- /* N */ qk_head_size,
616
- /* K */ v_head_size,
617
- /* ld_src */ v_head_size,
618
- /* ld_dst */ v_head_size);
680
+ if constexpr (brgemm_supported()) {
681
+ pack_vnni2<scalar_t>(
682
+ /* dst */ curr_last_recurrent_state_pack_reduced,
683
+ /* src */ curr_last_recurrent_state_reduced,
684
+ /* N */ qk_head_size,
685
+ /* K */ v_head_size,
686
+ /* ld_src */ v_head_size,
687
+ /* ld_dst */ v_head_size);
688
620
- // v_prime = k_cumdecay_i @ curr_last_recurrent_state: [chunk_size, EV]
621
- // k_cumdecay_i: [chunk_size, EK]
622
- // curr_last_recurrent_state: [EK, EV]
623
- at::native::cpublas::brgemm(
624
- /* M */ chunk_size,
625
- /* N */ v_head_size,
626
- /* K */ qk_head_size,
627
- /* lda */ qk_head_size,
628
- /* ldb */ v_head_size,
629
- /* ldc */ v_head_size,
630
- /* add_C */ false,
631
- /* A */ k_cumdecay_i_reduced,
632
- /* B */ curr_last_recurrent_state_pack_reduced,
633
- /* C */ v_prime);
689
+ // v_prime = k_cumdecay_i @ curr_last_recurrent_state: [chunk_size, EV]
690
+ // k_cumdecay_i: [chunk_size, EK]
691
+ // curr_last_recurrent_state: [EK, EV]
692
+ at::native::cpublas::brgemm(
693
+ /* M */ chunk_size,
694
+ /* N */ v_head_size,
695
+ /* K */ qk_head_size,
696
+ /* lda */ qk_head_size,
697
+ /* ldb */ v_head_size,
698
+ /* ldc */ v_head_size,
699
+ /* add_C */ false,
700
+ /* A */ k_cumdecay_i_reduced,
701
+ /* B */ curr_last_recurrent_state_pack_reduced,
702
+ /* C */ v_prime);
703
+ } else {
704
+ blas_gemm(
705
+ at::native::TransposeType::NoTranspose,
706
+ at::native::TransposeType::NoTranspose,
707
+ v_head_size,
708
+ chunk_size,
709
+ qk_head_size,
710
+ 1.0f,
711
+ curr_last_recurrent_state_reduced,
712
+ v_head_size,
713
+ k_cumdecay_i_reduced,
714
+ qk_head_size,
715
+ 0.0f,
716
+ v_prime,
717
+ v_head_size);
718
+ }
719
720
// v_new = v_prime = v_i - v_prime
721
// v_i: [chunk_size, EV]
@@ -663,41 +748,75 @@ void chunk_gated_delta_rule_kernel_impl(
748
}
749
// attn_inter = qg @ curr_last_recurrent_state: [chunk_size, EV]
750
// curr_last_recurrent_state: [EK, EV]
666
- at::native::cpublas::brgemm(
667
- /* M */ chunk_size,
668
- /* N */ v_head_size,
669
- /* K */ qk_head_size,
670
- /* lda */ qk_head_size,
671
- /* ldb */ v_head_size,
672
- /* ldc */ v_head_size,
673
- /* add_C */ false,
674
- /* A */ qg,
675
- /* B */ curr_last_recurrent_state_pack_reduced,
676
- /* C */ attn_inter);
751
+ if constexpr (brgemm_supported()) {
752
+ at::native::cpublas::brgemm(
753
+ /* M */ chunk_size,
754
+ /* N */ v_head_size,
755
+ /* K */ qk_head_size,
756
+ /* lda */ qk_head_size,
757
+ /* ldb */ v_head_size,
758
+ /* ldc */ v_head_size,
759
+ /* add_C */ false,
760
+ /* A */ qg,
761
+ /* B */ curr_last_recurrent_state_pack_reduced,
762
+ /* C */ attn_inter);
763
+ } else {
764
+ blas_gemm(
765
+ at::native::TransposeType::NoTranspose,
766
+ at::native::TransposeType::NoTranspose,
767
+ v_head_size,
768
+ chunk_size,
769
+ qk_head_size,
770
+ 1.0f,
771
+ curr_last_recurrent_state_reduced,
772
+ v_head_size,
773
+ qg,
774
+ qk_head_size,
775
+ 0.0f,
776
+ attn_inter,
777
+ v_head_size);
778
+ }
779
780
// core_attn_out[:, :, i] = attn_inter + attn_i @ v_new
781
// pack for v_prime
680
- pack_vnni2<scalar_t>(
681
- /* dst */ v_prime_pack_reduced,
682
- /* src */ v_prime_reduced,
683
- /* N */ chunk_size,
684
- /* K */ v_head_size,
685
- /* ld_src */ v_head_size,
686
- /* ld_dst */ v_head_size);
687
- // attn_inter = attn_inter + attn_i @ v_new: [chunk_size, EV]
688
- // attn_i: [chunk_size, chunk_size]
689
- // v_new: [chunk_size, EV]
690
- at::native::cpublas::brgemm(
691
- /* M */ chunk_size,
692
- /* N */ v_head_size,
693
- /* K */ chunk_size,
694
- /* lda */ chunk_size,
695
- /* ldb */ v_head_size,
696
- /* ldc */ v_head_size,
697
- /* add_C */ true,
698
- /* A */ attn_i_reduced,
699
- /* B */ v_prime_pack_reduced,
700
- /* C */ attn_inter);
782
+ if constexpr (brgemm_supported()) {
783
+ pack_vnni2<scalar_t>(
784
+ /* dst */ v_prime_pack_reduced,
785
+ /* src */ v_prime_reduced,
786
+ /* N */ chunk_size,
787
+ /* K */ v_head_size,
788
+ /* ld_src */ v_head_size,
789
+ /* ld_dst */ v_head_size);
790
+ // attn_inter = attn_inter + attn_i @ v_new: [chunk_size, EV]
791
+ // attn_i: [chunk_size, chunk_size]
792
+ // v_new: [chunk_size, EV]
793
+ at::native::cpublas::brgemm(
794
+ /* M */ chunk_size,
795
+ /* N */ v_head_size,
796
+ /* K */ chunk_size,
797
+ /* lda */ chunk_size,
798
+ /* ldb */ v_head_size,
799
+ /* ldc */ v_head_size,
800
+ /* add_C */ true,
801
+ /* A */ attn_i_reduced,
802
+ /* B */ v_prime_pack_reduced,
803
+ /* C */ attn_inter);
804
+ } else {
805
+ blas_gemm(
806
+ at::native::TransposeType::NoTranspose,
807
+ at::native::TransposeType::NoTranspose,
808
+ v_head_size,
809
+ chunk_size,
810
+ chunk_size,
811
+ 1.0f,
812
+ v_prime_reduced,
813
+ v_head_size,
814
+ attn_i_reduced,
815
+ chunk_size,
816
+ 1.0f,
817
+ attn_inter,
818
+ v_head_size);
819
+ }
820
821
// core_attn_out[:, :, i] = attn_inter
822
for (int64_t m = 0; m < chunk_size; m++) {
@@ -762,17 +881,34 @@ void chunk_gated_delta_rule_kernel_impl(
881
/* ld_dst */ chunk_size);
882
// kgv = kg.transpose(-1, -2) @ v_new
883
// v_new: [chunk_size, EV]
765
- at::native::cpublas::brgemm(
766
- /* M */ qk_head_size,
767
- /* N */ v_head_size,
768
- /* K */ chunk_size,
769
- /* lda */ chunk_size,
770
- /* ldb */ v_head_size,
771
- /* ldc */ v_head_size,
772
- /* add_C */ false,
773
- /* A */ kg_transpose,
774
- /* B */ v_prime_pack_reduced,
775
- /* C */ kgv);
884
+ if constexpr (brgemm_supported()) {
885
+ at::native::cpublas::brgemm(
886
+ /* M */ qk_head_size,
887
+ /* N */ v_head_size,
888
+ /* K */ chunk_size,
889
+ /* lda */ chunk_size,
890
+ /* ldb */ v_head_size,
891
+ /* ldc */ v_head_size,
892
+ /* add_C */ false,
893
+ /* A */ kg_transpose,
894
+ /* B */ v_prime_pack_reduced,
895
+ /* C */ kgv);
896
+ } else {
897
+ blas_gemm(
898
+ at::native::TransposeType::NoTranspose,
899
+ at::native::TransposeType::NoTranspose,
900
+ v_head_size,
901
+ qk_head_size,
902
+ chunk_size,
903
+ 1.0f,
904
+ v_prime_reduced,
905
+ v_head_size,
906
+ kg_transpose,
907
+ chunk_size,
908
+ 0.0f,
909
+ kgv,
910
+ v_head_size);
911
+ }
912
// last_recurrent_state = 1) + 2)
913
for (int64_t m = 0; m < qk_head_size; m++) {
914
at::vec::map2<float>(
@@ -921,7 +1057,8 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
1057
float k_scale = use_qk_l2norm_in_kernel ? qk_scale_buf[k_scale_offset] : 1.0f;
1058
int64_t v_offset = si * v_strideS + bi * v_strideB + ni * v_strideH;
1059
int64_t o_offset = ((bi * seq_len + si) * v_num_heads + ni) * v_head_dim;
924
- float beta_val = 1 / (1 + std::exp(-b_ptr[ni]));
1060
+ // See: https://github.com/sgl-project/sglang/pull/26634
1061
+ float beta_val = 1 / (1 + std::exp(-b_ptr[bi * v_num_heads + ni]));
1062
fVec beta_vec = fVec(beta_val);
1063
int64_t dvi = 0;
1064
for (; dvi <= v_head_dim - VecSize; dvi += VecSize) {
csrc/cpu/sgl-kernels/gemm.h
+18
-7
@@ -4,9 +4,12 @@
4
// clang-format off
5
6
#pragma once
7
-#include <ATen/native/CPUBlas.h>
8
-
7
#include "common.h"
8
+#include "blas_gemm.h"
9
+
10
+#if defined(__AVX512F__) && defined(__AVX512BF16__) && defined(__AMX_BF16__)
11
+#define CPU_CAPABILITY_AVX512
12
+#endif
13
14
// amx-bf16
15
#define TILE_M 16
@@ -21,31 +24,39 @@ constexpr int block_size_n() {
24
return 2 * TILE_N;
25
}
26
27
+constexpr bool brgemm_supported() {
28
+#if defined(CPU_CAPABILITY_AVX512)
29
+ return true;
30
+#else
31
+ return false;
32
+#endif
33
+}
34
+
35
// define threshold using brgemm (intel AMX)
36
template <typename T>
37
inline bool can_use_brgemm(int M);
38
template <>
39
inline bool can_use_brgemm<at::BFloat16>(int M) {
29
- return M > 4;
40
+ return brgemm_supported() && M > 4;
41
}
42
template <>
43
inline bool can_use_brgemm<at::Half>(int M) {
33
- return true;
44
+ return brgemm_supported();
45
}
46
// this requires PyTorch 2.7 or above
47
template <>
48
inline bool can_use_brgemm<int8_t>(int M) {
38
- return M > 4;
49
+ return brgemm_supported() && M > 4;
50
}
51
52
template <>
53
inline bool can_use_brgemm<uint8_t>(int M) {
43
- return M > 4;
54
+ return brgemm_supported() && M > 4;
55
}
56
57
template <>
58
inline bool can_use_brgemm<at::Float8_e4m3fn>(int M) {
48
- return M > 4;
59
+ return brgemm_supported() && M > 4;
60
}
61
62
// work around compiler internal error
csrc/cpu/sgl-kernels/vec.h
+2
@@ -11,7 +11,9 @@
11
12
#include <ATen/cpu/vec/functional.h>
13
#include <ATen/cpu/vec/vec.h>
14
+#if defined(CPU_CAPABILITY_AVX512)
15
#include <immintrin.h>
16
+#endif
17
namespace {
18
19
using namespace at::vec;
csrc/cpu/torch_bindings.cpp
+19
-19
@@ -447,6 +447,25 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
447
"bool is_vnni) -> Tensor");
448
ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu);
449
450
+ // Adapted from sglang: casual_conv1d kernels
451
+ ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
452
+ ops.impl("causal_conv1d_weight_pack", torch::kCPU,
453
+ &causal_conv1d_weight_pack);
454
+ ops.def(
455
+ "causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "
456
+ "conv_states, Tensor? query_start_loc,"
457
+ "Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "
458
+ "int pad_slot_id, bool is_vnni) -> "
459
+ "Tensor");
460
+ ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);
461
+ ops.def(
462
+ "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
463
+ "weight, Tensor? bias, bool silu_activation,"
464
+ "Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
465
+ "bool is_vnni) -> Tensor");
466
+ ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
467
+#endif
468
+
469
// Adapted from sglang: GDN kernels
470
ops.def(
471
"chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, "
@@ -470,25 +489,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
489
"-> (Tensor, Tensor)");
490
ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu);
491
473
- // Adapted from sglang: casual_conv1d kernels
474
- ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
475
- ops.impl("causal_conv1d_weight_pack", torch::kCPU,
476
- &causal_conv1d_weight_pack);
477
- ops.def(
478
- "causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "
479
- "conv_states, Tensor? query_start_loc,"
480
- "Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "
481
- "int pad_slot_id, bool is_vnni) -> "
482
- "Tensor");
483
- ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);
484
- ops.def(
485
- "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
486
- "weight, Tensor? bias, bool silu_activation,"
487
- "Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
488
- "bool is_vnni) -> Tensor");
489
- ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
490
-#endif
491
-
492
// CPU attention kernels
493
ops.def(
494
"get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, "
tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
new
+314
@@ -0,0 +1,314 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import functools
5
+
6
+import pytest
7
+import torch
8
+import torch.nn.functional as F
9
+
10
+import vllm._custom_ops as ops
11
+from vllm.platforms import current_platform
12
+from vllm.utils.torch_utils import set_random_seed
13
+
14
+if not current_platform.is_cpu():
15
+ pytest.skip("skipping CPU-only tests", allow_module_level=True)
16
+
17
+set_random_seed(12345)
18
+
19
+NUM_HEADS = [
20
+ (2, 4),
21
+ (4, 4),
22
+]
23
+HEAD_DIMS = [
24
+ (32, 32),
25
+ (64, 32),
26
+]
27
+CHUNK_SIZE = 64
28
+PREFILL_SEQ_LENS = [
29
+ [1],
30
+ [1, 2, 3],
31
+ [CHUNK_SIZE - 1],
32
+ [CHUNK_SIZE],
33
+ [CHUNK_SIZE + 1],
34
+ [CHUNK_SIZE - 1, CHUNK_SIZE, CHUNK_SIZE + 1],
35
+ [2 * CHUNK_SIZE - 1, 2 * CHUNK_SIZE, 2 * CHUNK_SIZE + 1],
36
+ [4 * CHUNK_SIZE + 17],
37
+]
38
+DECODE_BATCH_SIZES = [1, 3, 5]
39
+
40
+
41
+@functools.lru_cache(maxsize=128, typed=False)
42
+def tensor_cache(
43
+ elem_num: int,
44
+ dtype: torch.dtype,
45
+) -> torch.Tensor:
46
+ tensor = torch.rand(elem_num, dtype=dtype)
47
+ return tensor
48
+
49
+
50
+def ref_l2norm(
51
+ x: torch.Tensor,
52
+ dim: int = -1,
53
+ eps: float = 1e-5,
54
+) -> torch.Tensor:
55
+ inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
56
+ return x * inv_norm
57
+
58
+
59
+def ref_gdn_gating(
60
+ A_log: torch.Tensor,
61
+ a: torch.Tensor,
62
+ b: torch.Tensor,
63
+ dt_bias: torch.Tensor,
64
+) -> tuple[torch.Tensor, torch.Tensor]:
65
+ softplus_x = F.softplus(a.float() + dt_bias.float(), beta=1.0, threshold=20.0)
66
+ g = -torch.exp(A_log.float()) * softplus_x
67
+ beta = torch.sigmoid(b.float()).to(dtype=b.dtype)
68
+ return g, beta
69
+
70
+
71
+def ref_gated_delta_rule(
72
+ query: torch.Tensor,
73
+ key: torch.Tensor,
74
+ value: torch.Tensor,
75
+ a: torch.Tensor,
76
+ b: torch.Tensor,
77
+ A_log: torch.Tensor,
78
+ dt_bias: torch.Tensor,
79
+ initial_state: torch.Tensor,
80
+ cu_seqlens: torch.Tensor,
81
+ use_qk_l2norm_in_kernel: bool = False,
82
+) -> tuple[torch.Tensor, torch.Tensor]:
83
+ g, beta = ref_gdn_gating(A_log, a, b, dt_bias)
84
+ out = torch.empty_like(value)
85
+ final_state = torch.empty_like(initial_state)
86
+
87
+ for seq_idx in range(cu_seqlens.numel() - 1):
88
+ begin = int(cu_seqlens[seq_idx].item())
89
+ end = int(cu_seqlens[seq_idx + 1].item())
90
+ q_seq = query[:, begin:end]
91
+ k_seq = key[:, begin:end]
92
+ v_seq = value[:, begin:end]
93
+ g_seq = g[begin:end].unsqueeze(0)
94
+ beta_seq = beta[begin:end].unsqueeze(0)
95
+ initial_dtype = q_seq.dtype
96
+
97
+ if use_qk_l2norm_in_kernel:
98
+ q_seq = ref_l2norm(q_seq, dim=-1)
99
+ k_seq = ref_l2norm(k_seq, dim=-1)
100
+
101
+ if q_seq.shape[2] != v_seq.shape[2]:
102
+ repeat_factor = v_seq.shape[2] // q_seq.shape[2]
103
+ q_seq = q_seq.repeat_interleave(repeat_factor, dim=2)
104
+ k_seq = k_seq.repeat_interleave(repeat_factor, dim=2)
105
+
106
+ q_seq, k_seq, v_seq, beta_seq, g_seq = [
107
+ x.transpose(1, 2).contiguous().to(torch.float32)
108
+ for x in (q_seq, k_seq, v_seq, beta_seq, g_seq)
109
+ ]
110
+
111
+ batch_size, num_heads, seq_len, head_dim = q_seq.shape
112
+ v_head_dim = v_seq.shape[-1]
113
+ q_seq = q_seq * (1 / (head_dim**0.5))
114
+ out_seq = torch.empty(
115
+ batch_size,
116
+ num_heads,
117
+ seq_len,
118
+ v_head_dim,
119
+ dtype=v_seq.dtype,
120
+ )
121
+ state = initial_state[seq_idx : seq_idx + 1].to(v_seq)
122
+
123
+ for token_idx in range(seq_len):
124
+ q_t = q_seq[:, :, token_idx]
125
+ k_t = k_seq[:, :, token_idx]
126
+ v_t = v_seq[:, :, token_idx]
127
+ g_t = g_seq[:, :, token_idx].exp().unsqueeze(-1).unsqueeze(-1)
128
+ beta_t = beta_seq[:, :, token_idx].unsqueeze(-1)
129
+
130
+ state = state * g_t
131
+ kv_mem = (state * k_t.unsqueeze(-2)).sum(dim=-1)
132
+ delta = (v_t - kv_mem) * beta_t
133
+ state = state + delta.unsqueeze(-1) * k_t.unsqueeze(-2)
134
+ out_seq[:, :, token_idx] = (state * q_t.unsqueeze(-2)).sum(dim=-1)
135
+
136
+ out[:, begin:end] = out_seq.transpose(1, 2).contiguous().to(initial_dtype)
137
+ final_state[seq_idx] = state.squeeze(0)
138
+
139
+ return out, final_state
140
+
141
+
142
+def gdn_inputs(
143
+ num_tokens: int,
144
+ num_heads: tuple[int, int],
145
+ head_dims: tuple[int, int],
146
+) -> tuple[torch.Tensor, ...]:
147
+ num_qk_heads, num_v_heads = num_heads
148
+ head_dim, v_head_dim = head_dims
149
+ q_shape = (1, num_tokens, num_qk_heads, head_dim)
150
+ q_numel = num_tokens * num_qk_heads * head_dim
151
+ q = tensor_cache(q_numel, torch.bfloat16).view(q_shape)
152
+ k = tensor_cache(q_numel, torch.bfloat16).view(q_shape)
153
+
154
+ v_shape = (1, num_tokens, num_v_heads, v_head_dim)
155
+ v = tensor_cache(num_tokens * num_v_heads * v_head_dim, torch.bfloat16).view(
156
+ v_shape
157
+ )
158
+
159
+ gate_shape = (num_tokens, num_v_heads)
160
+ gate_numel = num_tokens * num_v_heads
161
+ a = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape)
162
+ b = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape)
163
+ A_log = tensor_cache(num_v_heads, torch.float32)
164
+ dt_bias = tensor_cache(num_v_heads, torch.bfloat16)
165
+ return q, k, v, a, b, A_log, dt_bias
166
+
167
+
168
+@pytest.mark.parametrize("num_tokens", [1, 9])
169
+@pytest.mark.parametrize("num_v_heads", [4, 8])
170
+@torch.inference_mode()
171
+def test_fused_gdn_gating_cpu(
172
+ num_tokens: int,
173
+ num_v_heads: int,
174
+) -> None:
175
+ gate_shape = (num_tokens, num_v_heads)
176
+ gate_numel = num_tokens * num_v_heads
177
+ a = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape)
178
+ b = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape)
179
+ A_log = tensor_cache(num_v_heads, torch.float32)
180
+ dt_bias = tensor_cache(num_v_heads, torch.bfloat16)
181
+
182
+ g_ref, beta_ref = ref_gdn_gating(A_log, a, b, dt_bias)
183
+ g, beta = ops.fused_gdn_gating_cpu(A_log, a, b, dt_bias)
184
+
185
+ torch.testing.assert_close(g, g_ref.unsqueeze(0), atol=1e-4, rtol=1e-4)
186
+ torch.testing.assert_close(
187
+ beta.float(), beta_ref.unsqueeze(0).float(), atol=5e-3, rtol=5e-3
188
+ )
189
+
190
+
191
+# decode path
192
+@pytest.mark.parametrize("batch_size", DECODE_BATCH_SIZES)
193
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
194
+@pytest.mark.parametrize("head_dims", HEAD_DIMS)
195
+@torch.inference_mode()
196
+def test_fused_sigmoid_gating_delta_rule_update_cpu(
197
+ batch_size: int,
198
+ num_heads: tuple[int, int],
199
+ head_dims: tuple[int, int],
200
+) -> None:
201
+ q, k, v, a, b, A_log, dt_bias = gdn_inputs(
202
+ num_tokens=batch_size,
203
+ num_heads=num_heads,
204
+ head_dims=head_dims,
205
+ )
206
+ _, num_v_heads = num_heads
207
+ head_dim, v_head_dim = head_dims
208
+ state_indices = torch.arange(batch_size, dtype=torch.int32)
209
+ cu_seqlens = torch.arange(batch_size + 1, dtype=torch.int32)
210
+ state_shape = (batch_size, num_v_heads, head_dim, v_head_dim)
211
+ state = tensor_cache(
212
+ batch_size * num_v_heads * head_dim * v_head_dim, torch.float32
213
+ ).view(state_shape)
214
+ state_ref = state[state_indices].transpose(-1, -2).contiguous()
215
+
216
+ out_ref, final_state_ref = ref_gated_delta_rule(
217
+ query=q,
218
+ key=k,
219
+ value=v,
220
+ a=a,
221
+ b=b,
222
+ A_log=A_log,
223
+ dt_bias=dt_bias,
224
+ initial_state=state_ref,
225
+ cu_seqlens=cu_seqlens,
226
+ use_qk_l2norm_in_kernel=True,
227
+ )
228
+ out_ref = out_ref.transpose(0, 1).contiguous()
229
+
230
+ state_out = state.clone()
231
+ out = ops.fused_sigmoid_gating_delta_rule_update_cpu(
232
+ A_log=A_log,
233
+ dt_bias=dt_bias,
234
+ q=q,
235
+ k=k,
236
+ v=v,
237
+ a=a,
238
+ b=b,
239
+ initial_state_source=state_out,
240
+ initial_state_indices=state_indices,
241
+ cu_seqlens=cu_seqlens,
242
+ use_qk_l2norm_in_kernel=True,
243
+ )
244
+
245
+ torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2)
246
+ torch.testing.assert_close(
247
+ state_out[state_indices].transpose(-1, -2),
248
+ final_state_ref,
249
+ atol=1e-2,
250
+ rtol=1e-2,
251
+ )
252
+
253
+
254
+# prefill path
255
+@pytest.mark.parametrize("seq_lens", PREFILL_SEQ_LENS)
256
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
257
+@pytest.mark.parametrize("head_dims", HEAD_DIMS)
258
+@torch.inference_mode()
259
+def test_chunk_gated_delta_rule_cpu(
260
+ seq_lens: list[int],
261
+ num_heads: tuple[int, int],
262
+ head_dims: tuple[int, int],
263
+) -> None:
264
+ total_tokens = sum(seq_lens)
265
+ q, k, v, a, b, A_log, dt_bias = gdn_inputs(
266
+ num_tokens=total_tokens,
267
+ num_heads=num_heads,
268
+ head_dims=head_dims,
269
+ )
270
+ _, num_v_heads = num_heads
271
+ head_dim, v_head_dim = head_dims
272
+ cu_seqlens = torch.tensor(
273
+ [0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32
274
+ )
275
+ initial_state_shape = (len(seq_lens), num_v_heads, head_dim, v_head_dim)
276
+ initial_state = tensor_cache(
277
+ len(seq_lens) * num_v_heads * head_dim * v_head_dim, torch.float32
278
+ ).view(initial_state_shape)
279
+ initial_state_ref = initial_state.transpose(-1, -2).contiguous()
280
+
281
+ out_ref, final_state_ref = ref_gated_delta_rule(
282
+ query=q,
283
+ key=k,
284
+ value=v,
285
+ a=a,
286
+ b=b,
287
+ A_log=A_log,
288
+ dt_bias=dt_bias,
289
+ initial_state=initial_state_ref,
290
+ cu_seqlens=cu_seqlens,
291
+ use_qk_l2norm_in_kernel=True,
292
+ )
293
+
294
+ g, beta = ref_gdn_gating(A_log, a, b, dt_bias)
295
+ out, final_state = ops.chunk_gated_delta_rule_cpu(
296
+ query=q,
297
+ key=k,
298
+ value=v,
299
+ g=g.unsqueeze(0),
300
+ beta=beta.unsqueeze(0),
301
+ initial_state=initial_state,
302
+ output_final_state=True,
303
+ cu_seqlens=cu_seqlens,
304
+ head_first=False,
305
+ use_qk_l2norm_in_kernel=True,
306
+ )
307
+
308
+ torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2)
309
+ torch.testing.assert_close(
310
+ final_state.transpose(-1, -2),
311
+ final_state_ref,
312
+ atol=1e-2,
313
+ rtol=1e-2,
314
+ )
vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py
+76
-192
@@ -12,11 +12,6 @@ from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
12
causal_conv1d_torch,
13
causal_conv1d_update_torch,
14
)
15
-from vllm.model_executor.layers.mamba.ops.cpu.recurrent_gated_delta_rule import (
16
- chunk_gated_delta_rule,
17
- gdn_gating,
18
- recurrent_gated_delta_rule,
19
-)
15
from vllm.utils.torch_utils import (
16
LayerNameType,
17
_resolve_layer_name,
@@ -55,179 +50,34 @@ def cpu_gdn_attention_core(
50
attn_metadata_i.spec_sequence_masks is None
51
and attn_metadata_i.num_accepted_tokens is None
52
), "speculative decode not supported in CPU GDN attention."
58
-
59
- if torch.cpu._is_amx_tile_supported():
60
- return cpu_gdn_attention_core_amx(
61
- mixed_qkv,
62
- b,
63
- a,
64
- core_attn_out,
65
- attn_metadata_i,
66
- layer,
67
- )
53
+ assert mixed_qkv.dtype == torch.bfloat16, "CPU GDN attention requires BF16."
54
55
state_indices_tensor = attn_metadata_i.non_spec_state_indices_tensor
56
query_start_loc = attn_metadata_i.non_spec_query_start_loc
57
assert state_indices_tensor is not None
58
assert query_start_loc is not None
59
74
- # [num_allocated_slots, conv_dim, kernel - 1]
75
- conv_state = layer.kv_cache[0]
76
- if not is_conv_state_dim_first():
77
- conv_state = conv_state.transpose(-1, -2)
78
-
79
- # [num_allocated_slots, num_v_heads / tp_size, v_dim, k_dim]
80
- ssm_state = layer.kv_cache[1]
81
-
82
- num_decodes = attn_metadata_i.num_decodes
83
- num_decode_tokens = attn_metadata_i.num_decode_tokens
84
- num_prefills = attn_metadata_i.num_prefills
85
- num_prefill_tokens = attn_metadata_i.num_prefill_tokens
86
-
87
- conv_weights = layer.conv1d.weight.view(
88
- layer.conv1d.weight.size(0), layer.conv1d.weight.size(2)
89
- )
90
-
91
- # all decode requests (batched)
92
- if num_decodes > 0:
93
- decode_mixed_qkv = mixed_qkv[:num_decode_tokens]
94
- decode_b = b[:num_decode_tokens]
95
- decode_a = a[:num_decode_tokens]
96
- decode_state_indices = state_indices_tensor[:num_decodes]
97
- decode_conv_state = conv_state[decode_state_indices].contiguous()
98
-
99
- decode_mixed_qkv = causal_conv1d_update_torch(
100
- # [B, dim] -> [B, dim, 1]
101
- x=decode_mixed_qkv.unsqueeze(-1),
102
- conv_state=decode_conv_state,
103
- weight=conv_weights,
104
- bias=layer.conv1d.bias,
105
- activation=layer.activation,
106
- ).squeeze(-1)
107
- conv_state[decode_state_indices] = decode_conv_state
108
-
109
- query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv)
110
-
111
- # [1, L, H, D] -> [B, 1, H, D] for batched decode
112
- query = query.transpose(0, 1).contiguous()
113
- key = key.transpose(0, 1).contiguous()
114
- value = value.transpose(0, 1).contiguous()
115
-
116
- g, beta_output = gdn_gating(
117
- A_log=layer.A_log,
118
- a=decode_a,
119
- b=decode_b,
120
- dt_bias=layer.dt_bias,
121
- )
122
- if g.ndim == 2:
123
- g = g.unsqueeze(1)
124
- beta_output = beta_output.unsqueeze(1)
125
-
126
- initial_state = ssm_state[decode_state_indices].contiguous()
127
- attn_out, last_recurrent_state = recurrent_gated_delta_rule(
128
- query=query,
129
- key=key,
130
- value=value,
131
- g=g,
132
- beta=beta_output,
133
- initial_state=initial_state,
134
- scale=None,
135
- use_qk_l2norm_in_kernel=True,
136
- )
137
- ssm_state[decode_state_indices] = last_recurrent_state.to(
138
- ssm_state.dtype
139
- ).contiguous()
140
- core_attn_out[:num_decode_tokens] = attn_out.squeeze(1)
141
-
142
- # all prefill requests: (varlen) currently naively loops over sequences
143
- if num_prefills > 0:
144
- has_initial_state = attn_metadata_i.has_initial_state
145
- assert has_initial_state is not None
146
-
147
- prefill_token_start = num_decode_tokens
148
- prefill_token_end = prefill_token_start + num_prefill_tokens
149
- prefill_mixed_qkv = mixed_qkv[prefill_token_start:prefill_token_end]
150
- prefill_b = b[prefill_token_start:prefill_token_end]
151
- prefill_a = a[prefill_token_start:prefill_token_end]
152
- prefill_state_indices = state_indices_tensor[
153
- num_decodes : num_decodes + num_prefills
154
- ]
155
- prefill_query_start_loc = (
156
- query_start_loc[num_decodes : num_decodes + num_prefills + 1]
157
- - num_decode_tokens
158
- )
159
- prefill_has_initial_state = has_initial_state[
160
- num_decodes : num_decodes + num_prefills
161
- ]
162
-
163
- prefill_mixed_qkv = causal_conv1d_torch(
164
- x=prefill_mixed_qkv.transpose(0, 1),
165
- weight=conv_weights,
166
- bias=layer.conv1d.bias,
167
- conv_states=conv_state,
168
- query_start_loc=prefill_query_start_loc,
169
- cache_indices=prefill_state_indices,
170
- has_initial_state=prefill_has_initial_state,
171
- activation=layer.activation,
172
- ).transpose(0, 1)
173
-
174
- query, key, value = layer.rearrange_mixed_qkv(prefill_mixed_qkv)
175
- g, beta = gdn_gating(layer.A_log, prefill_a, prefill_b, layer.dt_bias)
176
- if g.ndim == 2:
177
- g = g.unsqueeze(0)
178
- beta = beta.unsqueeze(0)
179
-
180
- initial_state = ssm_state[prefill_state_indices].contiguous()
181
- initial_state[~prefill_has_initial_state, ...] = 0
182
- attn_out, last_recurrent_state = chunk_gated_delta_rule(
183
- q=query,
184
- k=key,
185
- v=value,
186
- g=g,
187
- beta=beta,
188
- scale=None,
189
- initial_state=initial_state,
190
- cu_seqlens=prefill_query_start_loc,
191
- use_qk_l2norm_in_kernel=True,
192
- )
193
- ssm_state[prefill_state_indices] = last_recurrent_state.to(ssm_state.dtype)
194
- core_attn_out[prefill_token_start:prefill_token_end] = attn_out.squeeze(0)
195
-
60
+ is_amx = torch.cpu._is_amx_tile_supported()
61
197
-def cpu_gdn_attention_core_fake(
198
- mixed_qkv: torch.Tensor,
199
- b: torch.Tensor,
200
- a: torch.Tensor,
201
- core_attn_out: torch.Tensor,
202
- layer_name: LayerNameType,
203
-) -> None:
204
- """Fake implementation for torch.compile."""
205
- return
206
-
207
-
208
-def cpu_gdn_attention_core_amx(
209
- mixed_qkv: torch.Tensor,
210
- b: torch.Tensor,
211
- a: torch.Tensor,
212
- core_attn_out: torch.Tensor,
213
- attn_metadata_i: GDNAttentionMetadata,
214
- layer: torch.nn.Module,
215
-):
216
- state_indices_tensor = attn_metadata_i.non_spec_state_indices_tensor
217
- query_start_loc = attn_metadata_i.non_spec_query_start_loc
218
- assert state_indices_tensor is not None
219
- assert query_start_loc is not None
220
-
221
- # [num_allocated_slots, kernel - 1, conv_dim]
62
conv_state = layer.kv_cache[0]
223
- if is_conv_state_dim_first():
224
- raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.")
225
- # reshape to [num_allocated_slots, conv_dim, kernel - 1]
226
- conv_state_t = conv_state.transpose(1, 2)
63
+ if is_amx:
64
+ # AMX causal conv requires [num_allocated_slots, kernel - 1, conv_dim].
65
+ if is_conv_state_dim_first():
66
+ raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.")
67
+ conv_state = conv_state.transpose(1, 2)
68
+ else:
69
+ if not is_conv_state_dim_first():
70
+ conv_state = conv_state.transpose(-1, -2)
71
+ conv_weights = layer.conv1d.weight.view(
72
+ layer.conv1d.weight.size(0), layer.conv1d.weight.size(2)
73
+ )
74
75
# [num_allocated_slots, num_v_heads / tp_size, v_dim, k_dim]
76
ssm_state = layer.kv_cache[1]
230
- # rehape to [num_allocated_slots, num_v_heads / tp_size, k_dim, v_dim]
77
+ mixed_qkv = mixed_qkv.contiguous()
78
+ a = a.contiguous()
79
+ b = b.contiguous()
80
+
81
num_allocated_slots, head_num, v_dim, k_dim = ssm_state.size()
82
ssm_state = ssm_state.view(
83
num_allocated_slots,
@@ -236,32 +86,42 @@ def cpu_gdn_attention_core_amx(
86
v_dim,
87
)
88
239
- mixed_qkv = mixed_qkv.contiguous()
240
- a = a.contiguous()
241
- b = b.contiguous()
242
-
89
num_decodes = attn_metadata_i.num_decodes
90
num_decode_tokens = attn_metadata_i.num_decode_tokens
91
num_prefills = attn_metadata_i.num_prefills
92
num_prefill_tokens = attn_metadata_i.num_prefill_tokens
93
94
+ # all decode requests (batched)
95
if num_decodes > 0:
96
decode_mixed_qkv = mixed_qkv[:num_decode_tokens]
97
decode_b = b[:num_decode_tokens]
98
decode_a = a[:num_decode_tokens]
99
decode_state_indices = state_indices_tensor[:num_decodes]
253
-
254
- decode_mixed_qkv = ops.causal_conv1d_update_cpu(
255
- x=decode_mixed_qkv,
256
- conv_states=conv_state_t,
257
- weight=layer.conv1d.weight,
258
- bias=layer.conv1d.bias,
259
- silu_activation=layer.activation == "silu",
260
- conv_state_indices=decode_state_indices,
261
- is_vnni=True,
262
- )
100
+ if is_amx:
101
+ decode_mixed_qkv = ops.causal_conv1d_update_cpu(
102
+ x=decode_mixed_qkv,
103
+ conv_states=conv_state,
104
+ weight=layer.conv1d.weight,
105
+ bias=layer.conv1d.bias,
106
+ silu_activation=layer.activation == "silu",
107
+ conv_state_indices=decode_state_indices,
108
+ is_vnni=True,
109
+ )
110
+ else:
111
+ decode_conv_state = conv_state[decode_state_indices].contiguous()
112
+
113
+ decode_mixed_qkv = causal_conv1d_update_torch(
114
+ # [B, dim] -> [B, dim, 1]
115
+ x=decode_mixed_qkv.unsqueeze(-1),
116
+ conv_state=decode_conv_state,
117
+ weight=conv_weights,
118
+ bias=layer.conv1d.bias,
119
+ activation=layer.activation,
120
+ ).squeeze(-1)
121
+ conv_state[decode_state_indices] = decode_conv_state
122
123
query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv)
124
+
125
attn_out = ops.fused_sigmoid_gating_delta_rule_update_cpu(
126
A_log=layer.A_log,
127
dt_bias=layer.dt_bias,
@@ -277,6 +137,7 @@ def cpu_gdn_attention_core_amx(
137
)
138
core_attn_out[:num_decode_tokens] = attn_out.squeeze(1)
139
140
+ # all prefill requests: (varlen) currently naively loops over sequences
141
if num_prefills > 0:
142
has_initial_state = attn_metadata_i.has_initial_state
143
assert has_initial_state is not None
@@ -297,17 +158,29 @@ def cpu_gdn_attention_core_amx(
158
num_decodes : num_decodes + num_prefills
159
]
160
300
- prefill_mixed_qkv = ops.causal_conv1d_fwd_cpu(
301
- x=prefill_mixed_qkv.transpose(0, 1),
302
- weight=layer.conv1d.weight,
303
- bias=layer.conv1d.bias,
304
- conv_states=conv_state_t,
305
- query_start_loc=prefill_query_start_loc,
306
- cache_indices=prefill_state_indices,
307
- has_initial_state=prefill_has_initial_state,
308
- silu_activation=layer.activation == "silu",
309
- is_vnni=True,
310
- ).transpose(0, 1)
161
+ if is_amx:
162
+ prefill_mixed_qkv = ops.causal_conv1d_fwd_cpu(
163
+ x=prefill_mixed_qkv.transpose(0, 1),
164
+ weight=layer.conv1d.weight,
165
+ bias=layer.conv1d.bias,
166
+ conv_states=conv_state,
167
+ query_start_loc=prefill_query_start_loc,
168
+ cache_indices=prefill_state_indices,
169
+ has_initial_state=prefill_has_initial_state,
170
+ silu_activation=layer.activation == "silu",
171
+ is_vnni=True,
172
+ ).transpose(0, 1)
173
+ else:
174
+ prefill_mixed_qkv = causal_conv1d_torch(
175
+ x=prefill_mixed_qkv.transpose(0, 1),
176
+ weight=conv_weights,
177
+ bias=layer.conv1d.bias,
178
+ conv_states=conv_state,
179
+ query_start_loc=prefill_query_start_loc,
180
+ cache_indices=prefill_state_indices,
181
+ has_initial_state=prefill_has_initial_state,
182
+ activation=layer.activation,
183
+ ).transpose(0, 1)
184
185
query, key, value = layer.rearrange_mixed_qkv(prefill_mixed_qkv)
186
g, beta = ops.fused_gdn_gating_cpu(
@@ -334,6 +207,17 @@ def cpu_gdn_attention_core_amx(
207
core_attn_out[prefill_token_start:prefill_token_end] = attn_out.squeeze(0)
208
209
210
+def cpu_gdn_attention_core_fake(
211
+ mixed_qkv: torch.Tensor,
212
+ b: torch.Tensor,
213
+ a: torch.Tensor,
214
+ core_attn_out: torch.Tensor,
215
+ layer_name: LayerNameType,
216
+) -> None:
217
+ """Fake implementation for torch.compile."""
218
+ return
219
+
220
+
221
def register_cpu_gdn_attention_ops() -> None:
222
global _CPU_GDN_ATTENTION_OPS_REGISTERED
223
if _CPU_GDN_ATTENTION_OPS_REGISTERED:
vllm/model_executor/layers/mamba/ops/cpu/recurrent_gated_delta_rule.py
deleted
-223
@@ -1,223 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-
5
-import torch
6
-import torch.nn.functional as F
7
-
8
-
9
-def l2norm(
10
- x: torch.Tensor,
11
- dim: int = -1,
12
- eps: float = 1e-6,
13
-) -> torch.Tensor:
14
- inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
15
- return x * inv_norm
16
-
17
-
18
-def recurrent_gated_delta_rule(
19
- query: torch.Tensor,
20
- key: torch.Tensor,
21
- value: torch.Tensor,
22
- g: torch.Tensor,
23
- beta: torch.Tensor,
24
- initial_state: torch.Tensor,
25
- scale: float | None = None,
26
- use_qk_l2norm_in_kernel: bool = False,
27
-) -> tuple[torch.Tensor, torch.Tensor]:
28
- initial_dtype = query.dtype
29
- if use_qk_l2norm_in_kernel:
30
- query = l2norm(query, dim=-1, eps=1e-6)
31
- key = l2norm(key, dim=-1, eps=1e-6)
32
-
33
- if query.shape[2] != value.shape[2]:
34
- repeat_factor = value.shape[2] // query.shape[2]
35
- query = query.repeat_interleave(repeat_factor, dim=2)
36
- key = key.repeat_interleave(repeat_factor, dim=2)
37
-
38
- query, key, value, beta, g = [
39
- x.transpose(1, 2).contiguous().to(torch.float32)
40
- for x in (query, key, value, beta, g)
41
- ]
42
-
43
- batch_size, num_heads, sequence_length, _ = key.shape
44
- v_head_dim = value.shape[-1]
45
- if scale is None:
46
- scale = 1 / (query.shape[-1] ** 0.5)
47
- query = query * scale
48
-
49
- core_attn_out = torch.empty(
50
- batch_size,
51
- num_heads,
52
- sequence_length,
53
- v_head_dim,
54
- dtype=value.dtype,
55
- )
56
- last_recurrent_state = initial_state.to(value)
57
-
58
- for token_idx in range(sequence_length):
59
- q_t = query[:, :, token_idx]
60
- k_t = key[:, :, token_idx]
61
- v_t = value[:, :, token_idx]
62
- g_t = g[:, :, token_idx].exp().unsqueeze(-1).unsqueeze(-1)
63
- beta_t = beta[:, :, token_idx].unsqueeze(-1)
64
-
65
- last_recurrent_state = last_recurrent_state * g_t
66
- kv_mem = (last_recurrent_state * k_t.unsqueeze(-2)).sum(dim=-1)
67
- delta = (v_t - kv_mem) * beta_t
68
- last_recurrent_state = last_recurrent_state + delta.unsqueeze(
69
- -1
70
- ) * k_t.unsqueeze(-2)
71
- core_attn_out[:, :, token_idx] = (last_recurrent_state * q_t.unsqueeze(-2)).sum(
72
- dim=-1
73
- )
74
-
75
- core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
76
- return core_attn_out, last_recurrent_state
77
-
78
-
79
-def gdn_gating(
80
- A_log: torch.Tensor,
81
- a: torch.Tensor,
82
- b: torch.Tensor,
83
- dt_bias: torch.Tensor,
84
- beta: float = 1.0,
85
- threshold: float = 20.0,
86
-) -> tuple[torch.Tensor, torch.Tensor]:
87
- softplus_x = F.softplus(a.float() + dt_bias.float(), beta=beta, threshold=threshold)
88
- g = -torch.exp(A_log.float()) * softplus_x
89
- beta_output = torch.sigmoid(b.float()).to(dtype=b.dtype)
90
- return g, beta_output
91
-
92
-
93
-def chunk_gated_delta_rule(
94
- q: torch.Tensor,
95
- k: torch.Tensor,
96
- v: torch.Tensor,
97
- g: torch.Tensor,
98
- beta: torch.Tensor,
99
- *,
100
- initial_state: torch.Tensor,
101
- scale: float | None = None,
102
- cu_seqlens: torch.Tensor,
103
- use_qk_l2norm_in_kernel: bool = False,
104
-) -> tuple[torch.Tensor, torch.Tensor]:
105
- output = torch.empty_like(v)
106
- state_dtype = initial_state.dtype
107
- chunk_size = 128
108
- sequence_bounds = [
109
- (
110
- seq_idx,
111
- int(cu_seqlens[seq_idx].item()),
112
- int(cu_seqlens[seq_idx + 1].item()),
113
- )
114
- for seq_idx in range(len(cu_seqlens) - 1)
115
- ]
116
- chunk_eye = torch.eye(chunk_size, dtype=torch.float32)
117
- num_sequences = len(sequence_bounds)
118
- num_value_heads = v.shape[2]
119
- value_head_dim = v.shape[3]
120
- key_head_dim = k.shape[3]
121
- final_state = torch.empty(
122
- (num_sequences, num_value_heads, value_head_dim, key_head_dim),
123
- dtype=state_dtype,
124
- )
125
-
126
- for seq_idx, begin, end in sequence_bounds:
127
- q_seq = q[:, begin:end]
128
- k_seq = k[:, begin:end]
129
- v_seq = v[:, begin:end]
130
- g_seq = g[:, begin:end]
131
- beta_seq = beta[:, begin:end]
132
-
133
- initial_dtype = q_seq.dtype
134
- if use_qk_l2norm_in_kernel:
135
- q_seq = l2norm(q_seq, dim=-1, eps=1e-6)
136
- k_seq = l2norm(k_seq, dim=-1, eps=1e-6)
137
-
138
- num_qk_heads = q_seq.shape[2]
139
- num_value_heads = v_seq.shape[2]
140
- if num_qk_heads != num_value_heads:
141
- repeat_factor = num_value_heads // num_qk_heads
142
- q_seq = q_seq.repeat_interleave(repeat_factor, dim=2)
143
- k_seq = k_seq.repeat_interleave(repeat_factor, dim=2)
144
-
145
- q_seq, k_seq, v_seq, beta_seq, g_seq = [
146
- x.transpose(1, 2).contiguous().to(torch.float32)
147
- for x in (q_seq, k_seq, v_seq, beta_seq, g_seq)
148
- ]
149
- seq_batch_size, num_heads, seq_len, qk_head_dim = q_seq.shape
150
- value_head_dim = v_seq.shape[-1]
151
-
152
- if scale is None:
153
- scale = 1 / (qk_head_dim**0.5)
154
-
155
- q_seq = q_seq * scale
156
-
157
- seq_state = initial_state[seq_idx : seq_idx + 1].to(v_seq)
158
- seq_output = torch.empty(
159
- seq_batch_size,
160
- num_heads,
161
- seq_len,
162
- value_head_dim,
163
- dtype=v_seq.dtype,
164
- )
165
-
166
- for chunk_start in range(0, seq_len, chunk_size):
167
- chunk_end = min(chunk_start + chunk_size, seq_len)
168
- q_chunk = q_seq[:, :, chunk_start:chunk_end]
169
- k_chunk = k_seq[:, :, chunk_start:chunk_end]
170
- v_chunk = v_seq[:, :, chunk_start:chunk_end]
171
- beta_chunk = beta_seq[:, :, chunk_start:chunk_end]
172
- g_chunk = g_seq[:, :, chunk_start:chunk_end]
173
- chunk_len = chunk_end - chunk_start
174
-
175
- cum_g = g_chunk.cumsum(dim=-1)
176
- exp_cum_g = cum_g.exp()
177
- decay = (cum_g.unsqueeze(-1) - cum_g.unsqueeze(-2)).exp()
178
-
179
- interaction = (k_chunk * beta_chunk.unsqueeze(-1)) @ k_chunk.transpose(
180
- -1, -2
181
- )
182
- interaction = torch.tril(interaction * decay, diagonal=-1)
183
- system = interaction + chunk_eye[:chunk_len, :chunk_len]
184
-
185
- solved_values = torch.linalg.solve_triangular(
186
- system,
187
- v_chunk * beta_chunk.unsqueeze(-1),
188
- upper=False,
189
- )
190
- solved_keys = torch.linalg.solve_triangular(
191
- system,
192
- (k_chunk * beta_chunk.unsqueeze(-1)) * exp_cum_g.unsqueeze(-1),
193
- upper=False,
194
- )
195
-
196
- incoming_memory = torch.einsum("bhvk,bhck->bhcv", seq_state, solved_keys)
197
- transformed_values = solved_values - incoming_memory
198
-
199
- # Each chunk contributes both from the incoming recurrent state and
200
- # from its own in-chunk interactions.
201
- inter_chunk = torch.einsum(
202
- "bhvk,bhck->bhcv",
203
- seq_state,
204
- q_chunk * exp_cum_g.unsqueeze(-1),
205
- )
206
- intra_chunk = torch.tril((q_chunk @ k_chunk.transpose(-1, -2)) * decay)
207
- seq_output[:, :, chunk_start:chunk_end] = (
208
- inter_chunk + intra_chunk @ transformed_values
209
- )
210
-
211
- # Carry the recurrent state forward to the next chunk boundary.
212
- end_decay = (cum_g[:, :, -1:] - cum_g).exp().unsqueeze(-1)
213
- decayed_keys = k_chunk * end_decay
214
- seq_state = seq_state * exp_cum_g[:, :, -1, None, None] + torch.einsum(
215
- "bhcv,bhck->bhvk", transformed_values, decayed_keys
216
- )
217
-
218
- output[0, begin:end].copy_(
219
- seq_output.transpose(1, 2).contiguous().to(initial_dtype).squeeze(0)
220
- )
221
- final_state[seq_idx].copy_(seq_state.squeeze(0).to(state_dtype).contiguous())
222
-
223
- return output, final_state