minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Refining fused A GEMM dispatch (#31241)

Browse Source 
Co-authored-by: root <root@sgl-b300-inference.datacrunch.io>
Co-authored-by: Brayden Zhong <brayden@radixark.ai>
This commit is contained in:
Brayden Zhong
2026-07-16 16:57:01 +08:00
committed by GitHub
co-authored by root Brayden Zhong
parent a798a2aeea
commit e5f9804e26
5 changed files with 66 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/gemm/dsv3_fused_a_gemm.cuh
+41 -22
View File
@@ -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<bf16_t>(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<bf16_t>(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<uint32_t*>(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<uint32_t*>(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 <typename T, int kHdIn, int kHdOut, int kTileN, bool kUsePDL>
template <typename T, int kHdIn, int kHdOut, int kTileN, int kTileM, bool kUsePDL>
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 <int kHdIn, int kHdOut, bool kUsePDL>
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<bf16_t const*>(mat_b.data_ptr());
if (num_tokens <= 8) {
invokeFusedAGemm<bf16_t, kHdIn, kHdOut, 8, kUsePDL>(out_ptr, a_ptr, b_ptr, num_tokens, dev);
invokeFusedAGemm<bf16_t, kHdIn, kHdOut, 8, kTileM, kUsePDL>(out_ptr, a_ptr, b_ptr, num_tokens, dev);
} else {
invokeFusedAGemm<bf16_t, kHdIn, kHdOut, 16, kUsePDL>(out_ptr, a_ptr, b_ptr, num_tokens, dev);
invokeFusedAGemm<bf16_t, kHdIn, kHdOut, 16, kTileM, kUsePDL>(out_ptr, a_ptr, b_ptr, num_tokens, dev);
}
}
};
python/sglang/jit_kernel/fused_a_gemm.py
-10
View File
@@ -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]
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py
+2 -4
View File
@@ -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.
python/sglang/srt/models/deepseek_v2.py
+17
View File
@@ -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)