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 8, 03:45 UTC
[Model Bash] DeepSeek R1 BF16 Min Latency QKV A GEMM (0.5% E2E Speedup) (#34758)
Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>
Robert Shaw committed
Feb 18, 2026 at 10:42 UTC
6874638bc443ae99ee900072e0bb039fa5f7f0e7
7 files changed
+855
-3
CMakeLists.txt
+19
@@ -771,6 +771,25 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
771
endif()
772
endif()
773
774
+ # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
775
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
776
+ cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
777
+ else()
778
+ cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
779
+ endif()
780
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS)
781
+ set(DSV3_FUSED_A_GEMM_SRC "csrc/dsv3_fused_a_gemm.cu")
782
+ set_gencode_flags_for_srcs(
783
+ SRCS "${DSV3_FUSED_A_GEMM_SRC}"
784
+ CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}")
785
+ list(APPEND VLLM_EXT_SRC ${DSV3_FUSED_A_GEMM_SRC})
786
+ list(APPEND VLLM_GPU_FLAGS "-DENABLE_DSV3_FUSED_A_GEMM=1")
787
+ message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}")
788
+ else()
789
+ message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found "
790
+ "in CUDA target architectures.")
791
+ endif()
792
+
793
# moe_data.cu is used by all CUTLASS MoE kernels.
794
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
795
cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
csrc/dsv3_fused_a_gemm.cu
new
+747
@@ -0,0 +1,747 @@
1
+/*
2
+ * Adapted from
3
+ * https://github.com/sgl-project/sglang/blob/main/sgl-kernel/csrc/gemm/dsv3_fused_a_gemm.cu
4
+ * which was adapted from
5
+ * https://github.com/NVIDIA/TensorRT-LLM/blob/619709fc33bd5dc268f19d6a741fe7ed51c0f8f5/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3FusedAGemm.cu
6
+ *
7
+ * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved.
8
+ * Copyright (c) 2021, NAVER Corp. Authored by CLOVA.
9
+ *
10
+ * Licensed under the Apache License, Version 2.0 (the "License");
11
+ * you may not use this file except in compliance with the License.
12
+ * You may obtain a copy of the License at
13
+ *
14
+ * http://www.apache.org/licenses/LICENSE-2.0
15
+ *
16
+ * Unless required by applicable law or agreed to in writing, software
17
+ * distributed under the License is distributed on an "AS IS" BASIS,
18
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
19
+ * See the License for the specific language governing permissions and
20
+ * limitations under the License.
21
+ */
22
+
23
+#include <ATen/ATen.h>
24
+#include <ATen/cuda/CUDAContext.h>
25
+#include <cuda_bf16.h>
26
+#include <cuda_runtime.h>
27
+#include <torch/all.h>
28
+
29
+#include "core/registration.h"
30
+
31
+#include <cstdlib>
32
+#include <mutex>
33
+
34
+namespace {
35
+
36
+inline int getSMVersion() {
37
+ auto* props = at::cuda::getCurrentDeviceProperties();
38
+ return props->major * 10 + props->minor;
39
+}
40
+
41
+inline bool getEnvEnablePDL() {
42
+ static std::once_flag flag;
43
+ static bool enablePDL = false;
44
+ std::call_once(flag, [&]() {
45
+ if (getSMVersion() >= 90) {
46
+ char const* env = std::getenv("TRTLLM_ENABLE_PDL");
47
+ enablePDL = env && env[0] == '1' && env[1] == '\0';
48
+ }
49
+ });
50
+ return enablePDL;
51
+}
52
+
53
+} // namespace
54
+
55
+using bf16_t = __nv_bfloat16;
56
+
57
+__device__ void hmma_16_8_16_f32acc_bf16ab(float (&d_reg)[4],
58
+ const bf16_t (&a_reg)[8],
59
+ const bf16_t (&b_reg)[4],
60
+ float const (&c_reg)[4]) {
61
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
62
+ uint32_t a0 = *reinterpret_cast<uint32_t const*>(a_reg + 0);
63
+ uint32_t a1 = *reinterpret_cast<uint32_t const*>(a_reg + 2);
64
+ uint32_t a2 = *reinterpret_cast<uint32_t const*>(a_reg + 4);
65
+ uint32_t a3 = *reinterpret_cast<uint32_t const*>(a_reg + 6);
66
+ uint32_t b0 = *reinterpret_cast<uint32_t const*>(b_reg + 0);
67
+ uint32_t b1 = *reinterpret_cast<uint32_t const*>(b_reg + 2);
68
+ asm volatile(
69
+ "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
70
+ "{%0, %1, %2, %3},"
71
+ "{%4, %5, %6, %7},"
72
+ "{%8, %9},"
73
+ "{%10, %11, %12, %13};\n"
74
+ : "=f"(d_reg[0]), "=f"(d_reg[1]), "=f"(d_reg[2]), "=f"(d_reg[3])
75
+ : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "f"(d_reg[0]),
76
+ "f"(d_reg[1]), "f"(d_reg[2]), "f"(d_reg[3]));
77
+#endif
78
+}
79
+
80
+extern "C" {
81
+__device__ uint32_t __nvvm_get_smem_pointer(void*);
82
+}
83
+
84
+__device__ void ldgsts_128(void const* gPtr, void* sPtr, uint32_t pred) {
85
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
86
+ if (pred) {
87
+ uint32_t smemPtrAsUint32 = __nvvm_get_smem_pointer(sPtr);
88
+ asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2;\n" ::"r"(
89
+ smemPtrAsUint32),
90
+ "l"(gPtr), "n"(16));
91
+ }
92
+#endif
93
+}
94
+
95
+__device__ void ldsm_x4(void* smem_ptr, uint32_t* reg_ptr) {
96
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
97
+ asm volatile(
98
+ "ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n"
99
+ : "=r"(reg_ptr[0]), "=r"(reg_ptr[1]), "=r"(reg_ptr[2]), "=r"(reg_ptr[3])
100
+ : "r"(__nvvm_get_smem_pointer(smem_ptr)));
101
+#endif
102
+}
103
+
104
+template <class Type>
105
+__device__ int apply_swizzle_343_on_elem_row_col(int row_idx_, int col_idx_) {
106
+ uint32_t row_idx = *reinterpret_cast<uint32_t*>(&row_idx_);
107
+ uint32_t col_idx = *reinterpret_cast<uint32_t*>(&col_idx_);
108
+ row_idx = row_idx % 8;
109
+ row_idx = row_idx * (16 / sizeof(Type));
110
+ col_idx = col_idx ^ row_idx;
111
+ return *reinterpret_cast<int*>(&col_idx);
112
+}
113
+
114
+__device__ void initialize_barrier(
115
+ uint64_t* smem_barrier, // 64 bits user-manged barrier in smem
116
+ int thread_count =
117
+ 1) // Thread count expected to arrive/wait on this barrier
118
+{
119
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
120
+ uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
121
+ asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(smem_int_ptr),
122
+ "r"(thread_count));
123
+#endif
124
+}
125
+
126
+// Barrier wait
127
+__device__ void wait_barrier(
128
+ uint64_t* smem_barrier, // 64 bits user-manged barrier in smem
129
+ int phase_bit) // Current phase bit the barrier waiting to flip
130
+{
131
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
132
+ uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
133
+ asm volatile(
134
+ "{\n"
135
+ ".reg .pred P1;\n"
136
+ "LAB_WAIT:\n"
137
+ "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
138
+ "@P1 bra DONE;\n"
139
+ "bra LAB_WAIT;\n"
140
+ "DONE:\n"
141
+ "}\n" ::"r"(smem_int_ptr),
142
+ "r"(phase_bit));
143
+#endif
144
+}
145
+
146
+__device__ bool try_wait_barrier(uint64_t* smem_ptr, int phase_bit) {
147
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
148
+ uint32_t wait_complete;
149
+ uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_ptr);
150
+ asm volatile(
151
+ "{\n\t"
152
+ ".reg .pred P1; \n\t"
153
+ "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2; \n\t"
154
+ "selp.b32 %0, 1, 0, P1; \n\t"
155
+ "}"
156
+ : "=r"(wait_complete)
157
+ : "r"(smem_int_ptr), "r"(phase_bit));
158
+ return static_cast<bool>(wait_complete);
159
+#endif
160
+ return false;
161
+}
162
+
163
+// Barrier arrive
164
+__device__ void arrive_barrier(
165
+ uint64_t* smem_barrier) // 64 bits user-manged barrier in smem
166
+{
167
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
168
+ uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
169
+ asm volatile(
170
+ "{\n"
171
+ ".reg .b64 state; \n"
172
+ "mbarrier.arrive.shared::cta.b64 state, [%0];\n"
173
+ "}\n" ::"r"(smem_int_ptr));
174
+#endif
175
+}
176
+
177
+__device__ void ldgsts_arrive(uint64_t* smem_barrier) {
178
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
179
+ uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
180
+ asm volatile("cp.async.mbarrier.arrive.noinc.shared.b64 [%0];"
181
+ :
182
+ : "r"(smem_int_ptr));
183
+#endif
184
+}
185
+
186
+template <int gemm_k, int tile_m, int tile_k, int stage_cnt>
187
+struct GmemLoaderA {
188
+ static constexpr int elem_bytes = 2;
189
+ static constexpr int vec_bytes = 16;
190
+ static constexpr int vec_elems = vec_bytes / elem_bytes;
191
+ static constexpr int thread_cnt = 64;
192
+ static_assert((tile_m * tile_k) % (vec_elems * thread_cnt) == 0);
193
+ static constexpr int a_inst_cnt_per_iter =
194
+ (tile_m * tile_k) / (vec_elems * thread_cnt);
195
+ static_assert(gemm_k % tile_k == 0);
196
+ static constexpr int k_iter_cnt = gemm_k / tile_k;
197
+
198
+ // Extra params to keep the order of k reduction...
199
+ static constexpr int mma_warp_cnt = 4;
200
+ static constexpr int per_mma_warp_k = tile_k / mma_warp_cnt;
201
+ static constexpr int k_each_chunk = gemm_k / mma_warp_cnt;
202
+
203
+ private:
204
+ __device__ int k_project(int tile_k_idx) {
205
+ return (tile_k_idx / per_mma_warp_k * k_each_chunk) +
206
+ (tile_k_idx % per_mma_warp_k);
207
+ }
208
+
209
+ public:
210
+ __device__ GmemLoaderA(bf16_t const* gmem_a_local_, bf16_t* smem_a_,
211
+ uint64_t* smem_barrier_)
212
+ : gmem_a(gmem_a_local_),
213
+ smem_a(smem_a_),
214
+ smem_barrier(smem_barrier_),
215
+ local_tid(threadIdx.x % thread_cnt) {}
216
+
217
+ __device__ void prepare() {
218
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
219
+ // swizzle, that's what we want.
220
+ #pragma unroll
221
+ for (int i = 0; i < a_inst_cnt_per_iter; i++) {
222
+ int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
223
+ int m_idx = linear_idx / tile_k;
224
+ int k_idx = linear_idx % tile_k;
225
+ k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(m_idx, k_idx);
226
+ a_smem_offsets[i] = m_idx * tile_k + k_idx;
227
+ }
228
+#endif
229
+ }
230
+
231
+ __device__ void issue_mainloop() {
232
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
233
+ #pragma unroll 1
234
+ for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
235
+ if (need_wait) {
236
+ wait_barrier(smem_barrier + 1 + stage_idx * 2, phase_bit);
237
+ }
238
+ int next_stage_idx = stage_idx + 1;
239
+ int next_phase_bit =
240
+ next_stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
241
+ next_stage_idx = next_stage_idx == stage_cnt ? 0 : next_stage_idx;
242
+ if (loop_idx != k_iter_cnt - 1) {
243
+ need_wait = !try_wait_barrier(smem_barrier + 1 + next_stage_idx * 2,
244
+ next_phase_bit);
245
+ }
246
+
247
+ #pragma unroll
248
+ for (int i = 0; i < a_inst_cnt_per_iter; i++) {
249
+ int smem_offset = a_smem_offsets[i];
250
+ bf16_t* smem_ptr_this_iter =
251
+ smem_a + stage_idx * tile_m * tile_k + smem_offset;
252
+ int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
253
+ int m_idx = linear_idx / tile_k;
254
+ int k_idx = linear_idx % tile_k;
255
+ int gmem_offset = m_idx * gemm_k + k_project(k_idx);
256
+ bf16_t const* gmem_ptr_this_iter = gmem_a + gmem_offset;
257
+ ldgsts_128(gmem_ptr_this_iter, smem_ptr_this_iter, true);
258
+ }
259
+ ldgsts_arrive(smem_barrier + stage_idx * 2);
260
+
261
+ stage_idx = next_stage_idx;
262
+ phase_bit = next_phase_bit;
263
+ gmem_a += per_mma_warp_k;
264
+ }
265
+#endif
266
+ }
267
+
268
+ bf16_t const* gmem_a;
269
+ bf16_t* smem_a;
270
+ uint64_t* smem_barrier;
271
+ int local_tid;
272
+ int stage_idx = 0;
273
+ int phase_bit = 1;
274
+ bool need_wait = true;
275
+
276
+ // per smem_stage, store with swizzle information
277
+ int a_smem_offsets[a_inst_cnt_per_iter];
278
+};
279
+
280
+template <int gemm_k, int tile_n, int tile_k, int stage_cnt>
281
+struct GmemLoaderB {
282
+ static constexpr int elem_bytes = 2;
283
+ static constexpr int vec_bytes = 16;
284
+ static constexpr int vec_elems = vec_bytes / elem_bytes;
285
+ static constexpr int thread_cnt = 64;
286
+ static_assert((tile_n * tile_k) % (vec_elems * thread_cnt) == 0);
287
+ static constexpr int b_inst_cnt_per_iter =
288
+ (tile_n * tile_k) / (vec_elems * thread_cnt);
289
+ static_assert(gemm_k % tile_k == 0);
290
+ static constexpr int k_iter_cnt = gemm_k / tile_k;
291
+
292
+ // Extra params to keep the order of k reduction...
293
+ static constexpr int mma_warp_cnt = 4;
294
+ static constexpr int per_mma_warp_k = tile_k / mma_warp_cnt;
295
+ static constexpr int k_each_chunk = gemm_k / mma_warp_cnt;
296
+
297
+ private:
298
+ __device__ int k_project(int tile_k_idx) {
299
+ return (tile_k_idx / per_mma_warp_k * k_each_chunk) +
300
+ (tile_k_idx % per_mma_warp_k);
301
+ }
302
+
303
+ public:
304
+ __device__ GmemLoaderB(bf16_t const* gmem_b_local_, bf16_t* smem_b_,
305
+ uint64_t* smem_barrier_, int gemm_n_)
306
+ : gmem_b(gmem_b_local_),
307
+ smem_b(smem_b_),
308
+ smem_barrier(smem_barrier_),
309
+ gemm_n(gemm_n_),
310
+ local_tid(threadIdx.x % thread_cnt) {}
311
+
312
+ __device__ void prepare() {
313
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
314
+ // swizzle, that's what we want.
315
+ #pragma unroll
316
+ for (int i = 0; i < b_inst_cnt_per_iter; i++) {
317
+ int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
318
+ int n_idx = linear_idx / tile_k;
319
+ int k_idx = linear_idx % tile_k;
320
+ k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(n_idx, k_idx);
321
+ b_smem_offsets[i] = n_idx * tile_k + k_idx;
322
+ preds[i] = n_idx < gemm_n;
323
+ }
324
+#endif
325
+ }
326
+
327
+ __device__ void issue_mainloop() {
328
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
329
+ asm volatile("griddepcontrol.wait;");
330
+ #pragma unroll 1
331
+ for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
332
+ if (need_wait) {
333
+ wait_barrier(smem_barrier + 1 + stage_idx * 2, phase_bit);
334
+ }
335
+ int next_stage_idx = stage_idx + 1;
336
+ int next_phase_bit =
337
+ next_stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
338
+ next_stage_idx = next_stage_idx == stage_cnt ? 0 : next_stage_idx;
339
+ if (loop_idx != k_iter_cnt - 1) {
340
+ need_wait = !try_wait_barrier(smem_barrier + 1 + next_stage_idx * 2,
341
+ next_phase_bit);
342
+ }
343
+ #pragma unroll
344
+ for (int i = 0; i < b_inst_cnt_per_iter; i++) {
345
+ int smem_offset = b_smem_offsets[i];
346
+ bf16_t* smem_ptr_this_iter =
347
+ smem_b + stage_idx * tile_n * tile_k + smem_offset;
348
+ int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
349
+ int n_idx = linear_idx / tile_k;
350
+ int k_idx = linear_idx % tile_k;
351
+ int gmem_offset = n_idx * gemm_k + k_project(k_idx);
352
+ bf16_t const* gmem_ptr_this_iter = gmem_b + gmem_offset;
353
+ ldgsts_128(gmem_ptr_this_iter, smem_ptr_this_iter, preds[i]);
354
+ }
355
+ ldgsts_arrive(smem_barrier + stage_idx * 2);
356
+
357
+ stage_idx = next_stage_idx;
358
+ phase_bit = next_phase_bit;
359
+ gmem_b += per_mma_warp_k;
360
+ }
361
+#endif
362
+ }
363
+
364
+ bf16_t const* gmem_b;
365
+ bf16_t* smem_b;
366
+ uint64_t* smem_barrier;
367
+ int gemm_n;
368
+ int local_tid;
369
+ int stage_idx = 0;
370
+ int phase_bit = 1;
371
+ bool need_wait = true;
372
+
373
+ // per smem_stage, store with swizzle information
374
+ int b_smem_offsets[b_inst_cnt_per_iter];
375
+ uint32_t preds[b_inst_cnt_per_iter];
376
+};
377
+
378
+template <int gemm_m, int gemm_k, int tile_m, int tile_n, int tile_k,
379
+ int stage_cnt>
380
+struct MmaComputer {
381
+ static constexpr int elem_bytes = 2;
382
+ static constexpr int thread_cnt = 128;
383
+ static_assert(gemm_k % tile_k == 0);
384
+ static_assert(tile_k % (thread_cnt / 32) == 0);
385
+ static constexpr int per_warp_tile_k = tile_k / (thread_cnt / 32);
386
+ static constexpr int k_iter_cnt = gemm_k / tile_k;
387
+ static constexpr int k_phase_cnt = per_warp_tile_k / 16;
388
+ static constexpr int m_iter_cnt = (tile_m + 15) / 16;
389
+ static constexpr int n_iter_cnt =
390
+ (tile_n + 7) /
391
+ 8; // Possible to have non-1 n_iter_cnt for ab_swap m16 case.
392
+ static_assert(m_iter_cnt == 1);
393
+ static_assert(n_iter_cnt == 1 || n_iter_cnt == 2);
394
+
395
+ __device__ MmaComputer(bf16_t* gmem_c_local_, bf16_t* smem_a_,
396
+ bf16_t* smem_b_, uint64_t* smem_barrier_,
397
+ int warp_idx_, int gemm_n_)
398
+ : gmem_c(gmem_c_local_),
399
+ smem_a(smem_a_),
400
+ smem_b(smem_b_),
401
+ smem_barrier(smem_barrier_),
402
+ warp_idx(warp_idx_ - (thread_cnt / 32)),
403
+ gemm_n(gemm_n_) {}
404
+
405
+ private:
406
+ __device__ constexpr int internal_b_atom_func(int tid) {
407
+ if constexpr (tile_n < 8) {
408
+ return (tid % tile_n) + ((tid % 8) / tile_n * 0) + tid / 8 * 8 * tile_n;
409
+ } else {
410
+ return (tid % 8) + ((tid % 32) / 8 * (tile_n * 8));
411
+ }
412
+ }
413
+
414
+ public:
415
+ __device__ void prepare() {
416
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
417
+ #pragma unroll
418
+ for (int i = 0; i < k_phase_cnt; i++) {
419
+ int linear_idx = (lane_idx % 16) + (lane_idx / 16) * 128 + i * 256;
420
+ int m_idx = linear_idx % tile_m;
421
+ int k_idx = linear_idx / tile_m + warp_k_offset_in_tile_k;
422
+ k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(m_idx, k_idx);
423
+ a_smem_offsets[0][i] = m_idx * tile_k + k_idx;
424
+ }
425
+ #pragma unroll
426
+ for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
427
+ #pragma unroll
428
+ for (int i = 0; i < k_phase_cnt; i += 2) { // Special i+=2 for B.
429
+ int linear_idx =
430
+ internal_b_atom_func(lane_idx) + i * tile_n * 16 + n_iter_idx * 8;
431
+ int n_idx = linear_idx % tile_n;
432
+ int k_idx = linear_idx / tile_n + warp_k_offset_in_tile_k;
433
+ k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(n_idx, k_idx);
434
+ b_smem_offsets[n_iter_idx][i] = n_idx * tile_k + k_idx;
435
+ }
436
+ }
437
+#endif
438
+ }
439
+
440
+ __device__ void issue_mainloop() {
441
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
442
+ #pragma unroll 1
443
+ for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
444
+ wait_barrier(smem_barrier + 0 + stage_idx * 2, phase_bit);
445
+
446
+ #pragma unroll
447
+ for (int i = 0; i < k_phase_cnt; i++) {
448
+ int smem_offset = a_smem_offsets[0][i];
449
+ bf16_t* smem_ptr_this_iter =
450
+ smem_a + stage_idx * tile_m * tile_k + smem_offset;
451
+ ldsm_x4(smem_ptr_this_iter, reinterpret_cast<uint32_t*>(a_reg[0][i]));
452
+ }
453
+
454
+ #pragma unroll
455
+ for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
456
+ #pragma unroll
457
+ for (int i = 0; i < k_phase_cnt; i += 2) {
458
+ int smem_offset = b_smem_offsets[n_iter_idx][i];
459
+ bf16_t* smem_ptr_this_iter =
460
+ smem_b + stage_idx * tile_n * tile_k + smem_offset;
461
+ ldsm_x4(smem_ptr_this_iter,
462
+ reinterpret_cast<uint32_t*>(b_reg[n_iter_idx][i]));
463
+ }
464
+ }
465
+
466
+ #pragma unroll
467
+ for (int k_iter_idx = 0; k_iter_idx < k_phase_cnt; k_iter_idx++) {
468
+ #pragma unroll
469
+ for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
470
+ hmma_16_8_16_f32acc_bf16ab(
471
+ acc_reg[0][n_iter_idx], a_reg[0][k_iter_idx],
472
+ b_reg[n_iter_idx][k_iter_idx], acc_reg[0][n_iter_idx]);
473
+ }
474
+ }
475
+ ::arrive_barrier(smem_barrier + 1 + stage_idx * 2);
476
+ stage_idx += 1;
477
+ phase_bit = stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
478
+ stage_idx = stage_idx == stage_cnt ? 0 : stage_idx;
479
+ }
480
+#endif
481
+ }
482
+
483
+ __device__ void epi() {
484
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
485
+ asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt));
486
+ // reorganize the acc_reg
487
+ constexpr int thread_m = 2;
488
+ constexpr int thread_n = 2 * n_iter_cnt;
489
+ constexpr int cta_mma_n = n_iter_cnt * 8;
490
+ float acc_reg_reorg[thread_m][thread_n];
491
+
492
+ for (int i = 0; i < thread_m; i++) {
493
+ for (int j = 0; j < thread_n; j++) {
494
+ acc_reg_reorg[i][j] = acc_reg[0][j / 2][(j % 2) + (i * 2)];
495
+ }
496
+ }
497
+
498
+ // 4 x cosize(smem_c_layout)
499
+ float* smem_c = reinterpret_cast<float*>(smem_a);
500
+ // coord -> index
501
+ auto smem_c_index_func = [&](int m_idx, int n_idx) {
502
+ int group_rows = 32 / cta_mma_n;
503
+ int group_cnt = 2;
504
+ return (m_idx % group_rows * cta_mma_n) +
505
+ (m_idx / group_rows * (32 + group_cnt)) + n_idx;
506
+ };
507
+ constexpr int cosize_smem_c = ((tile_m * cta_mma_n) / 32) * (32 + 2);
508
+
509
+ // This should be optimized to STS.64 but can not be STS.128 due to the bank
510
+ // index.
511
+ #pragma unroll
512
+ for (int m_idx_thread = 0; m_idx_thread < thread_m; m_idx_thread++) {
513
+ #pragma unroll
514
+ for (int n_idx_thread = 0; n_idx_thread < thread_n; n_idx_thread++) {
515
+ int m_idx = (lane_idx / 4) + m_idx_thread * 8;
516
+ int n_idx =
517
+ ((lane_idx % 4) * 2) + (n_idx_thread % 2) + (n_idx_thread / 2) * 8;
518
+ smem_c[cosize_smem_c * warp_idx + smem_c_index_func(m_idx, n_idx)] =
519
+ acc_reg_reorg[m_idx_thread][n_idx_thread];
520
+ }
521
+ }
522
+ asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt));
523
+
524
+ if (warp_idx == 0) {
525
+ constexpr int final_acc_reg_cnt = (tile_m * tile_n + 31) / 32;
526
+ float acc_final[final_acc_reg_cnt]{};
527
+
528
+ #pragma unroll
529
+ for (int reg_idx = 0; reg_idx < final_acc_reg_cnt; reg_idx++) {
530
+ int linear_idx = reg_idx * 32 + lane_idx;
531
+ int m_idx = linear_idx % tile_m;
532
+ int n_idx = linear_idx / tile_m;
533
+ acc_final[reg_idx] +=
534
+ smem_c[smem_c_index_func(m_idx, n_idx) + 0 * cosize_smem_c] +
535
+ smem_c[smem_c_index_func(m_idx, n_idx) + 1 * cosize_smem_c] +
536
+ smem_c[smem_c_index_func(m_idx, n_idx) + 2 * cosize_smem_c] +
537
+ smem_c[smem_c_index_func(m_idx, n_idx) + 3 * cosize_smem_c];
538
+ }
539
+
540
+ #pragma unroll
541
+ for (int reg_idx = 0; reg_idx < final_acc_reg_cnt; reg_idx++) {
542
+ int linear_idx = reg_idx * 32 + lane_idx;
543
+ int m_idx = linear_idx % tile_m;
544
+ int n_idx = linear_idx / tile_m;
545
+ if (m_idx < tile_m && n_idx < gemm_n) {
546
+ gmem_c[n_idx * gemm_m + m_idx] = acc_final[reg_idx];
547
+ }
548
+ }
549
+ }
550
+#endif
551
+ }
552
+
553
+ bf16_t* gmem_c;
554
+ bf16_t* smem_a;
555
+ bf16_t* smem_b;
556
+ uint64_t* smem_barrier;
557
+ int warp_idx;
558
+ int gemm_n;
559
+ int stage_idx = 0;
560
+ int phase_bit = 0;
561
+ int lane_idx = threadIdx.x % 32;
562
+ int warp_k_offset_in_tile_k = warp_idx * per_warp_tile_k;
563
+
564
+ int a_smem_offsets[m_iter_cnt][k_phase_cnt];
565
+ int b_smem_offsets[n_iter_cnt][k_phase_cnt];
566
+
567
+ bf16_t a_reg[m_iter_cnt][k_phase_cnt][8];
568
+ bf16_t b_reg[n_iter_cnt][k_phase_cnt][4];
569
+ float acc_reg[m_iter_cnt][n_iter_cnt][4]{};
570
+};
571
+
572
+// AB swapped, kernel is k-major, k-major, m-major
573
+template <int batch_size, int gemm_m, int gemm_k, int tile_m, int tile_n,
574
+ int tile_k, int stage_cnt>
575
+__global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
576
+ bf16_t* output, bf16_t const* mat_a, bf16_t const* mat_b, int gemm_n) {
577
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
578
+ constexpr int load_thread_cnt = 128;
579
+ constexpr int compute_thread_cnt = 128;
580
+ constexpr int thread_cnt = load_thread_cnt + compute_thread_cnt;
581
+ (void)thread_cnt;
582
+ static_assert(gemm_m % 16 == 0);
583
+ static_assert(gemm_k % tile_k == 0);
584
+ static_assert(gemm_m % tile_m == 0);
585
+ static_assert(
586
+ tile_k == 128 || tile_k == 256 || tile_k == 512 ||
587
+ tile_k == 1024); // tile_k must be larger than 64 since 4 warp splitK.
588
+ static_assert(tile_m == 16);
589
+ constexpr int g2s_vec_bytes = 16;
590
+ constexpr int a_elem_bytes = 2;
591
+ constexpr int b_elem_bytes = 2;
592
+ static_assert((tile_m * a_elem_bytes + tile_n * b_elem_bytes) * tile_k *
593
+ stage_cnt <=
594
+ 225 * 1024);
595
+ static_assert((tile_m * tile_k * a_elem_bytes) %
596
+ (load_thread_cnt * g2s_vec_bytes) ==
597
+ 0);
598
+ static_assert((tile_n * tile_k * b_elem_bytes) %
599
+ (load_thread_cnt * g2s_vec_bytes) ==
600
+ 0);
601
+
602
+ extern __shared__ char smem[];
603
+ uint64_t* smem_barrier = reinterpret_cast<uint64_t*>(
604
+ smem); // producer,consumer; producer,consumer; ...
605
+ bf16_t* smem_a = reinterpret_cast<bf16_t*>(smem + (stage_cnt * 8 * 2 + 1024) /
606
+ 1024 * 1024);
607
+ bf16_t* smem_b = smem_a + tile_m * tile_k * stage_cnt;
608
+
609
+ int cta_m_idx = tile_m * blockIdx.x;
610
+ int cta_n_idx = tile_n * blockIdx.y;
611
+ bf16_t const* gmem_a_local = mat_a + cta_m_idx * gemm_k;
612
+ bf16_t const* gmem_b_local = mat_b + cta_n_idx * gemm_k;
613
+ bf16_t* gmem_c_local = output + cta_n_idx * gemm_m + cta_m_idx;
614
+
615
+ int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
616
+
617
+ if (warp_idx == 4) {
618
+ for (int i = 0; i < stage_cnt; i++) {
619
+ initialize_barrier(smem_barrier + i * 2 + 0,
620
+ load_thread_cnt); // producer
621
+ initialize_barrier(smem_barrier + i * 2 + 1,
622
+ compute_thread_cnt); // consumer
623
+ }
624
+ }
625
+ __syncthreads();
626
+
627
+ if (warp_idx < 2) {
628
+ GmemLoaderA<gemm_k, tile_m, tile_k, stage_cnt> a_loader(
629
+ gmem_a_local, smem_a, smem_barrier);
630
+ a_loader.prepare();
631
+ a_loader.issue_mainloop();
632
+ } else if (warp_idx < 4) {
633
+ GmemLoaderB<gemm_k, tile_n, tile_k, stage_cnt> b_loader(
634
+ gmem_b_local, smem_b, smem_barrier, gemm_n);
635
+ b_loader.prepare();
636
+ b_loader.issue_mainloop();
637
+ } else {
638
+ MmaComputer<gemm_m, gemm_k, tile_m, tile_n, tile_k, stage_cnt> mma_computer(
639
+ gmem_c_local, smem_a, smem_b, smem_barrier, warp_idx, gemm_n);
640
+ mma_computer.prepare();
641
+ mma_computer.issue_mainloop();
642
+ mma_computer.epi();
643
+ }
644
+ asm volatile("griddepcontrol.launch_dependents;");
645
+#endif
646
+}
647
+
648
+template <typename T, int kHdIn, int kHdOut, int kTileN>
649
+void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
650
+ cudaStream_t const stream) {
651
+ constexpr int gemm_m = kHdOut; // 2112
652
+ int const gemm_n = num_tokens; // 1-16
653
+ constexpr int gemm_k = kHdIn; // 7168
654
+ constexpr int batch_size = 1;
655
+ std::swap(mat_a, mat_b);
656
+ constexpr int tile_m = 16;
657
+ constexpr int tile_n = kTileN; // 8 or 16
658
+ constexpr int tile_k = std::max(256, 1024 / tile_n); // 256
659
+ constexpr int max_stage_cnt =
660
+ 1024 * 192 / ((tile_m + tile_n) * tile_k * sizeof(bf16_t));
661
+ constexpr int k_iter_cnt = gemm_k / tile_k;
662
+ constexpr int stage_cnt =
663
+ k_iter_cnt > max_stage_cnt ? max_stage_cnt : k_iter_cnt;
664
+ int cta_m_cnt = gemm_m / tile_m;
665
+ int cta_n_cnt = (gemm_n + tile_n - 1) / tile_n;
666
+ constexpr int barrier_bytes = (stage_cnt * 16 + 1023) / 1024 * 1024;
667
+ constexpr int smem_bytes =
668
+ ((tile_m * 2 + tile_n * 2) * tile_k * stage_cnt + barrier_bytes + 1023) /
669
+ 1024 * 1024;
670
+
671
+ dim3 grid(cta_m_cnt, cta_n_cnt, 1);
672
+ dim3 block_size(256);
673
+ cudaLaunchConfig_t config;
674
+ config.gridDim = grid;
675
+ config.blockDim = block_size;
676
+ config.dynamicSmemBytes = smem_bytes;
677
+ config.stream = stream;
678
+ cudaLaunchAttribute attrs[1];
679
+ attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
680
+ attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
681
+ config.numAttrs = 1;
682
+ config.attrs = attrs;
683
+ if (smem_bytes >= (48 * 1024)) {
684
+ cudaFuncSetAttribute(fused_a_gemm_kernel<batch_size, gemm_m, gemm_k, tile_m,
685
+ tile_n, tile_k, stage_cnt>,
686
+ cudaFuncAttributeMaxDynamicSharedMemorySize,
687
+ smem_bytes);
688
+ }
689
+ cudaLaunchKernelEx(&config,
690
+ fused_a_gemm_kernel<batch_size, gemm_m, gemm_k, tile_m,
691
+ tile_n, tile_k, stage_cnt>,
692
+ output, mat_a, mat_b, gemm_n);
693
+}
694
+
695
+template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 8>(
696
+ __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
697
+ cudaStream_t);
698
+
699
+template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>(
700
+ __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
701
+ cudaStream_t);
702
+
703
+void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
704
+ torch::Tensor const& mat_b) {
705
+ TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2);
706
+ int const num_tokens = mat_a.size(0);
707
+ int const hd_in = mat_a.size(1);
708
+ int const hd_out = mat_b.size(1);
709
+
710
+ constexpr int kHdIn = 7168;
711
+ constexpr int kHdOut = 2112;
712
+ TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
713
+ "required 1 <= mat_a.shape[0] <= 16")
714
+ TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168")
715
+ TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112")
716
+ TORCH_CHECK(output.size(0) == num_tokens,
717
+ "required output.shape[0] == mat_a.shape[0]")
718
+ TORCH_CHECK(output.size(1) == hd_out,
719
+ "required output.shape[1] == mat_b.shape[1]")
720
+
721
+ TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor");
722
+ TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor");
723
+ TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor");
724
+
725
+ TORCH_CHECK(mat_a.scalar_type() == torch::kBFloat16 &&
726
+ mat_b.scalar_type() == torch::kBFloat16,
727
+ "Only BFloat16 input dtype is supported")
728
+ TORCH_CHECK(output.scalar_type() == torch::kBFloat16,
729
+ "Only BFloat16 output dtype is supported")
730
+
731
+ TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
732
+
733
+ auto stream = at::cuda::getCurrentCUDAStream(mat_a.get_device());
734
+ if (num_tokens <= 8) {
735
+ invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>(
736
+ reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
737
+ reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
738
+ reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
739
+ stream);
740
+ } else {
741
+ invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 16>(
742
+ reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
743
+ reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
744
+ reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
745
+ stream);
746
+ }
747
+}
csrc/ops.h
+5
@@ -410,3 +410,8 @@ void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
410
int64_t quant_level, bool cast_bf2half = false);
411
int64_t qr_max_size();
412
#endif
413
+
414
+#ifndef USE_ROCM
415
+void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
416
+ torch::Tensor const& mat_b);
417
+#endif
\ No newline at end of file
csrc/torch_bindings.cpp
+5
@@ -239,6 +239,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
239
240
// Quantization ops
241
#ifndef USE_ROCM
242
+ // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
243
+ ops.def(
244
+ "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
245
+ ops.impl("dsv3_fused_a_gemm", torch::kCUDA, &dsv3_fused_a_gemm);
246
+
247
// Quantized GEMM for AWQ.
248
ops.def(
249
"awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, "
vllm/_custom_ops.py
+18
@@ -2789,6 +2789,24 @@ def sm100_cutlass_mla_get_workspace_size(
2789
)
2790
2791
2792
+def dsv3_fused_a_gemm(
2793
+ output: torch.Tensor,
2794
+ mat_a: torch.Tensor,
2795
+ mat_b: torch.Tensor,
2796
+) -> None:
2797
+ """DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
2798
+
2799
+ Computes output = mat_a @ mat_b.T where:
2800
+ mat_a: [num_tokens, 7168] row-major bf16 (hidden states)
2801
+ mat_b: [7168, 2112] column-major bf16 (weight transposed)
2802
+ output: [num_tokens, 2112] row-major bf16
2803
+
2804
+ Optimized for the DeepSeek V2/V3 QKV A-projection at small batch sizes.
2805
+ Requires SM 9.0+ (Hopper).
2806
+ """
2807
+ torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b)
2808
+
2809
+
2810
if hasattr(torch.ops._C, "weight_packed_linear"):
2811
2812
@register_fake("_C::weight_packed_linear")
vllm/model_executor/layers/mla.py
+1
@@ -129,6 +129,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
129
assert self.q_b_proj is not None, (
130
"q_b_proj is required when q_lora_rank is not None"
131
)
132
+
133
qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
134
q_c, kv_lora = qkv_lora.split(
135
[self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
vllm/model_executor/models/deepseek_v2.py
+60
-3
@@ -32,6 +32,7 @@ import torch
32
from torch import nn
33
from transformers import DeepseekV2Config, DeepseekV3Config
34
35
+import vllm._custom_ops as ops
36
from vllm._aiter_ops import rocm_aiter_ops
37
from vllm.compilation.decorators import support_torch_compile
38
from vllm.config import CacheConfig, ParallelConfig, VllmConfig, get_current_vllm_config
@@ -711,6 +712,64 @@ class Indexer(nn.Module):
712
return self.indexer_op(hidden_states, q_fp8, k, weights)
713
714
715
+class DeepSeekV2FusedQkvAProj(MergedColumnParallelLinear):
716
+ def __init__(
717
+ self,
718
+ input_size: int,
719
+ output_size: int,
720
+ quant_config: QuantizationConfig | None = None,
721
+ prefix: str = "",
722
+ ):
723
+ super().__init__(
724
+ input_size,
725
+ output_size,
726
+ bias=False,
727
+ quant_config=quant_config,
728
+ disable_tp=True,
729
+ prefix=f"{prefix}.kv_a_proj_with_mqa",
730
+ )
731
+
732
+ # Check if the DeepSeek V3 fused A GEMM kernel can be used.
733
+ # This kernel supports PDL and is optimized for low batch size.
734
+ self._use_min_latency_gemm = (
735
+ hasattr(self, "weight")
736
+ and self.weight.dtype == torch.bfloat16
737
+ and self.weight.shape[0] == 2112
738
+ and self.weight.shape[1] == 7168
739
+ and current_platform.is_cuda()
740
+ and (
741
+ current_platform.is_device_capability(90)
742
+ or current_platform.is_device_capability_family(100)
743
+ )
744
+ )
745
+
746
+ def forward(
747
+ self,
748
+ input_,
749
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.nn.Parameter | None]:
750
+ num_tokens = input_.shape[0]
751
+ if self._use_min_latency_gemm and (0 < num_tokens <= 16):
752
+ output = torch.empty(
753
+ num_tokens,
754
+ 2112,
755
+ dtype=torch.bfloat16,
756
+ device=input_.device,
757
+ )
758
+ ops.dsv3_fused_a_gemm(
759
+ output,
760
+ input_,
761
+ self.weight.T,
762
+ )
763
+ if not self.return_bias:
764
+ return output
765
+ output_bias = self.bias if self.skip_bias_add else None
766
+ return output, output_bias
767
+ else:
768
+ # Fallback to the standard forward method when
769
+ # the fused A GEMM kernel cannot be used.
770
+ return super().forward(input_)
771
+
772
+
773
class DeepseekV2MLAAttention(nn.Module):
774
"""
775
Main reference: DeepseekV2 paper, and FlashInfer Implementation
@@ -756,13 +815,11 @@ class DeepseekV2MLAAttention(nn.Module):
815
self.max_position_embeddings = max_position_embeddings
816
817
if self.q_lora_rank is not None:
759
- self.fused_qkv_a_proj = MergedColumnParallelLinear(
818
+ self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProj(
819
self.hidden_size,
820
[self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
762
- bias=False,
821
quant_config=quant_config,
822
prefix=f"{prefix}.fused_qkv_a_proj",
765
- disable_tp=True,
823
)
824
else:
825
self.kv_a_proj_with_mqa = ReplicatedLinear(