From e5f9804e26c461383126dc60e953cd315729d2b7 Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Thu, 16 Jul 2026 01:57:01 -0700 Subject: [PATCH] Refining fused A GEMM dispatch (#31241) Co-authored-by: root Co-authored-by: Brayden Zhong --- .../csrc/gemm/dsv3_fused_a_gemm.cuh | 63 ++++++++++++------- python/sglang/jit_kernel/fused_a_gemm.py | 10 --- .../attention_forward_methods/forward_mla.py | 6 +- python/sglang/srt/models/deepseek_v2.py | 17 +++++ test/registered/jit/test_dsv3_fused_a_gemm.py | 10 +-- 5 files changed, 66 insertions(+), 40 deletions(-) diff --git a/python/sglang/jit_kernel/csrc/gemm/dsv3_fused_a_gemm.cuh b/python/sglang/jit_kernel/csrc/gemm/dsv3_fused_a_gemm.cuh index 804e44cdf..33584fbd9 100644 --- a/python/sglang/jit_kernel/csrc/gemm/dsv3_fused_a_gemm.cuh +++ b/python/sglang/jit_kernel/csrc/gemm/dsv3_fused_a_gemm.cuh @@ -333,7 +333,7 @@ struct MmaComputer { static constexpr int k_phase_cnt = per_warp_tile_k / 16; static constexpr int m_iter_cnt = (tile_m + 15) / 16; static constexpr int n_iter_cnt = (tile_n + 7) / 8; - static_assert(m_iter_cnt == 1); + static_assert(m_iter_cnt == 1 || m_iter_cnt == 2); static_assert(n_iter_cnt == 1 || n_iter_cnt == 2); __device__ MmaComputer( @@ -358,12 +358,15 @@ struct MmaComputer { __device__ void prepare() { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 #pragma unroll - for (int i = 0; i < k_phase_cnt; i++) { - int linear_idx = (lane_idx % 16) + (lane_idx / 16) * 128 + i * 256; - int m_idx = linear_idx % tile_m; - int k_idx = linear_idx / tile_m + warp_k_offset_in_tile_k; - k_idx = apply_swizzle_343_on_elem_row_col(m_idx, k_idx); - a_smem_offsets[0][i] = m_idx * tile_k + k_idx; + for (int m = 0; m < m_iter_cnt; m++) { +#pragma unroll + for (int i = 0; i < k_phase_cnt; i++) { + int linear_idx = (lane_idx % 16) + (lane_idx / 16) * 128 + i * 256; + int m_idx = linear_idx % 16 + m * 16; + int k_idx = linear_idx / 16 + warp_k_offset_in_tile_k; + k_idx = apply_swizzle_343_on_elem_row_col(m_idx, k_idx); + a_smem_offsets[m][i] = m_idx * tile_k + k_idx; + } } #pragma unroll for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) { @@ -386,10 +389,13 @@ struct MmaComputer { wait_barrier(smem_barrier + 0 + stage_idx * 2, phase_bit); #pragma unroll - for (int i = 0; i < k_phase_cnt; i++) { - int smem_offset = a_smem_offsets[0][i]; - bf16_t* smem_ptr_this_iter = smem_a + stage_idx * tile_m * tile_k + smem_offset; - ldsm_x4(smem_ptr_this_iter, reinterpret_cast(a_reg[0][i])); + for (int m = 0; m < m_iter_cnt; m++) { +#pragma unroll + for (int i = 0; i < k_phase_cnt; i++) { + int smem_offset = a_smem_offsets[m][i]; + bf16_t* smem_ptr_this_iter = smem_a + stage_idx * tile_m * tile_k + smem_offset; + ldsm_x4(smem_ptr_this_iter, reinterpret_cast(a_reg[m][i])); + } } #pragma unroll @@ -406,8 +412,11 @@ struct MmaComputer { for (int k_iter_idx = 0; k_iter_idx < k_phase_cnt; k_iter_idx++) { #pragma unroll for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) { - hmma_16_8_16_f32acc_bf16ab( - acc_reg[0][n_iter_idx], a_reg[0][k_iter_idx], b_reg[n_iter_idx][k_iter_idx], acc_reg[0][n_iter_idx]); +#pragma unroll + for (int m = 0; m < m_iter_cnt; m++) { + hmma_16_8_16_f32acc_bf16ab( + acc_reg[m][n_iter_idx], a_reg[m][k_iter_idx], b_reg[n_iter_idx][k_iter_idx], acc_reg[m][n_iter_idx]); + } } } ::arrive_barrier(smem_barrier + 1 + stage_idx * 2); @@ -421,14 +430,14 @@ struct MmaComputer { __device__ void epi() { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt)); - constexpr int thread_m = 2; + constexpr int thread_m = 2 * m_iter_cnt; constexpr int thread_n = 2 * n_iter_cnt; constexpr int cta_mma_n = n_iter_cnt * 8; float acc_reg_reorg[thread_m][thread_n]; for (int i = 0; i < thread_m; i++) { for (int j = 0; j < thread_n; j++) { - acc_reg_reorg[i][j] = acc_reg[0][j / 2][(j % 2) + (i * 2)]; + acc_reg_reorg[i][j] = acc_reg[i / 2][j / 2][(j % 2) + (i % 2) * 2]; } } @@ -444,7 +453,7 @@ struct MmaComputer { for (int m_idx_thread = 0; m_idx_thread < thread_m; m_idx_thread++) { #pragma unroll for (int n_idx_thread = 0; n_idx_thread < thread_n; n_idx_thread++) { - int m_idx = (lane_idx / 4) + m_idx_thread * 8; + int m_idx = (lane_idx / 4) + (m_idx_thread % 2) * 8 + (m_idx_thread / 2) * 16; int n_idx = ((lane_idx % 4) * 2) + (n_idx_thread % 2) + (n_idx_thread / 2) * 8; smem_c[cosize_smem_c * warp_idx + smem_c_index_func(m_idx, n_idx)] = acc_reg_reorg[m_idx_thread][n_idx_thread]; } @@ -472,7 +481,7 @@ struct MmaComputer { int m_idx = linear_idx % tile_m; int n_idx = linear_idx / tile_m; if (m_idx < tile_m && n_idx < gemm_n) { - gmem_c[n_idx * gemm_m + m_idx] = acc_final[reg_idx]; + gmem_c[n_idx * gemm_m + m_idx] = __float2bfloat16(acc_final[reg_idx]); } } } @@ -510,7 +519,7 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel( static_assert(gemm_k % tile_k == 0); static_assert(gemm_m % tile_m == 0); static_assert(tile_k == 128 || tile_k == 256 || tile_k == 512 || tile_k == 1024); - static_assert(tile_m == 16); + static_assert(tile_m == 16 || tile_m == 32); constexpr int g2s_vec_bytes = 16; constexpr int a_elem_bytes = 2; constexpr int b_elem_bytes = 2; @@ -558,14 +567,14 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel( #endif } -template +template void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens, DLDevice device) { constexpr int gemm_m = kHdOut; // 2112 int const gemm_n = num_tokens; // 16 constexpr int gemm_k = kHdIn; // 7168 constexpr int batch_size = 1; std::swap(mat_a, mat_b); - constexpr int tile_m = 16; + constexpr int tile_m = kTileM; constexpr int tile_n = kTileN; // 8 or 16 constexpr int tile_k = std::max(256, 1024 / tile_n); // 256 #if defined(SGL_CUDA_ARCH) && SGL_CUDA_ARCH >= 1200 @@ -591,8 +600,18 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens, host::LaunchKernel(grid, block_size, device, smem_bytes).enable_pdl(kUsePDL)(kernel, output, mat_a, mat_b, gemm_n); } +// tile_m=32 halves the CTA count so these shapes fit in one SM wave instead of two. +constexpr int pick_tile_m(int hd_in, int hd_out) { + if (hd_out == 2624 && hd_in == 6144) return 32; + if (hd_out == 4096 && hd_in == 2048) return 32; + return 16; +} + template struct DSV3FusedAGemmKernel { + static constexpr int kTileM = pick_tile_m(kHdIn, kHdOut); + static_assert(kHdOut % kTileM == 0, "hd_out must be a multiple of tile_m"); + static void run(const tvm::ffi::TensorView mat_a, const tvm::ffi::TensorView mat_b, const tvm::ffi::TensorView output) { using namespace host; @@ -621,9 +640,9 @@ struct DSV3FusedAGemmKernel { auto* b_ptr = static_cast(mat_b.data_ptr()); if (num_tokens <= 8) { - invokeFusedAGemm(out_ptr, a_ptr, b_ptr, num_tokens, dev); + invokeFusedAGemm(out_ptr, a_ptr, b_ptr, num_tokens, dev); } else { - invokeFusedAGemm(out_ptr, a_ptr, b_ptr, num_tokens, dev); + invokeFusedAGemm(out_ptr, a_ptr, b_ptr, num_tokens, dev); } } }; diff --git a/python/sglang/jit_kernel/fused_a_gemm.py b/python/sglang/jit_kernel/fused_a_gemm.py index 7d893444a..f5e3446e3 100644 --- a/python/sglang/jit_kernel/fused_a_gemm.py +++ b/python/sglang/jit_kernel/fused_a_gemm.py @@ -16,7 +16,6 @@ from enum import Enum import torch -from sglang.srt.layers.quantization.unquant import get_bf16_gemm_backend from sglang.srt.utils.common import get_device_sm, is_cuda, is_sm120_supported @@ -52,19 +51,10 @@ def linear_with_fused_a_gemm( backend: "FusedAGemmBackend | str" = FusedAGemmBackend.AUTO, ) -> torch.Tensor: # LoRA reads weight.T directly, bypassing the adapter, so fall back when active. - cutedsl_backend = get_bf16_gemm_backend().is_cutedsl() - if cutedsl_backend: - from sglang.jit_kernel.cutedsl_bf16_gemm import use_cutedsl_bf16_gemm if ( not isinstance(hidden_states, tuple) and 1 <= hidden_states.shape[0] <= 16 and not getattr(layer, "set_lora", False) - and not ( - cutedsl_backend - and use_cutedsl_bf16_gemm( - hidden_states.shape[0], layer.weight.shape[0], layer.weight.shape[1] - ) - ) ): return dsv3_fused_a_gemm(hidden_states, layer.weight.T, backend=backend) return layer(hidden_states)[0] diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index da0af746f..c3ed46742 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -377,9 +377,7 @@ class DeepseekMLAForwardMixin: self.alt_stream.wait_stream(current_stream) with torch.cuda.stream(self.alt_stream): k_nope = k_nope.unsqueeze(1) - q = self.q_b_proj(q)[0].view( - -1, self.num_local_heads, self.qk_head_dim - ) + q = self.q_b_proj_forward(q) if self.should_run_indexer(prev_topk_indices): topk_indices = self.indexer( x=hidden_states, @@ -397,7 +395,7 @@ class DeepseekMLAForwardMixin: current_stream.wait_stream(self.alt_stream) else: k_nope = k_nope.unsqueeze(1) - q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) + q = self.q_b_proj_forward(q) # Hoist these above the DSA indexer split op so the indexer # and the composite bmm+attention split op are adjacent in FX. diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 73b7a4b05..c0087d34a 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -1787,6 +1787,14 @@ class DeepseekV2AttentionMLA( ) self.fused_a_gemm_backend = "auto" + self.has_q_b_proj = hasattr(self, "q_b_proj") + q_b_proj_verified_shapes = {(2048, 2048), (4096, 2048)} + self.use_min_latency_q_b_gemm = ( + self.has_q_b_proj + and tuple(self.q_b_proj.weight.shape) in q_b_proj_verified_shapes + and fused_a_gemm_weight_eligible(self.q_b_proj) + ) + self.init_mha_forward() self.init_mla_forward() self.init_mla_fused_rope_rocm_forward() @@ -1996,6 +2004,15 @@ class DeepseekV2AttentionMLA( ) return self.fused_qkv_a_proj_with_mqa(hidden_states)[0] + def q_b_proj_forward(self, q_lora: torch.Tensor) -> torch.Tensor: + if self.use_min_latency_q_b_gemm: + q = linear_with_fused_a_gemm( + self.q_b_proj, q_lora, backend=self.fused_a_gemm_backend + ) + else: + q = self.q_b_proj(q_lora)[0] + return q.view(-1, self.num_local_heads, self.qk_head_dim) + def rebuild_cp_kv_cache(self, latent_cache, forward_batch, k_nope, k_pe): # support allgather+rerrange latent_cache[..., : self.kv_lora_rank] = k_nope.squeeze(1) diff --git a/test/registered/jit/test_dsv3_fused_a_gemm.py b/test/registered/jit/test_dsv3_fused_a_gemm.py index c223b1d32..83287f0e0 100644 --- a/test/registered/jit/test_dsv3_fused_a_gemm.py +++ b/test/registered/jit/test_dsv3_fused_a_gemm.py @@ -12,10 +12,12 @@ from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large") -# hd_in must be a multiple of 256; 6144/7168 cover the real fused-A shapes. -HD_INS = [6144, 7168] -# hd_out must be a multiple of 16; 2112 and 2624 cover real fused-A variants. -HD_OUTS = [2112, 2624] +# hd_in must be a multiple of 256; 2048/6144/7168 cover the real fused-A shapes +# (2048 is q_b_proj TP4/TP8, 6144/7168 are qkv_a). +HD_INS = [2048, 6144, 7168] +# hd_out must be a multiple of 16; 2048/2112/2624/4096 cover real fused-A variants +# (2048/4096 are q_b_proj TP8/TP4, 2112/2624 are qkv_a). +HD_OUTS = [2048, 2112, 2624, 4096] NUM_TOKENS = get_ci_test_range(list(range(1, 17)), [1, 8, 16])