From e6c582866aeb9c5f52d03db6fa716258fce8df1f Mon Sep 17 00:00:00 2001 From: apinge Date: Tue, 4 Aug 2026 12:07:14 +0000 Subject: [PATCH 1/2] add workround for qkv_proj --- python/sglang/srt/speculative/dflash_utils.py | 25 +++++++------------ 1 file changed, 9 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index ede45007cd7f..dd1c36ad8eb9 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -534,26 +534,19 @@ def can_dflash_slice_qkv_weight(qkv_proj: Any) -> Tuple[bool, str]: return False, "qkv weight tensor is missing" weight = qkv_proj.weight if getattr(weight, "aiter_trans_weight", False): - return ( - False, - "qkv weight uses AITER shuffled layout; direct tensor slicing would bypass " - "the linear method", - ) - if weight.ndim != 2: + return False, "AITER transposed qkv weight layout is not sliceable" + if getattr(weight, "ndim", None) != 2: return False, f"qkv weight must be 2D, got shape={tuple(weight.shape)}" - expected_in_features = getattr(qkv_proj, "input_size", None) - if expected_in_features is None: - expected_in_features = getattr(qkv_proj, "hidden_size", None) - if ( - expected_in_features is not None - and int(weight.shape[1]) != int(expected_in_features) - ): + expected_shape = ( + int(getattr(qkv_proj, "output_size_per_partition")), + int(getattr(qkv_proj, "input_size")), + ) + if tuple(weight.shape) != expected_shape: return ( False, - "qkv weight input dim does not match hidden size; direct tensor slicing " - f"would be invalid (weight.shape={tuple(weight.shape)}, " - f"expected_in_features={int(expected_in_features)})", + "qkv weight layout is not sliceable: " + f"expected shape={expected_shape}, got shape={tuple(weight.shape)}", ) return True, "" From cebdc33667dd01e2f45bf3a2ab0ddcee3ba10e26 Mon Sep 17 00:00:00 2001 From: apinge Date: Tue, 4 Aug 2026 13:13:30 +0000 Subject: [PATCH 2/2] fix qkv_proj slice --- python/sglang/srt/models/dflash.py | 17 +++++++++ python/sglang/srt/speculative/dflash_utils.py | 37 ++++++++++++++++++- 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/models/dflash.py b/python/sglang/srt/models/dflash.py index 5f27b7fb5bd2..d567b534e7ab 100644 --- a/python/sglang/srt/models/dflash.py +++ b/python/sglang/srt/models/dflash.py @@ -22,6 +22,9 @@ RowParallelLinear, ) from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.layers.quantization.compressed_tensors.utils import ( + rocm_aiter_swizzle_hipb_unquantized_gemm, +) from sglang.srt.layers.radix_attention import AttentionType, RadixAttention from sglang.srt.layers.rotary_embedding import get_rope from sglang.srt.model_executor.forward_batch_info import ForwardBatch @@ -29,6 +32,7 @@ from sglang.srt.models.utils import apply_qk_norm from sglang.srt.runtime_context import get_parallel from sglang.srt.speculative.dflash_utils import ( + can_dflash_slice_aiter_trans_qkv_weight, can_dflash_slice_qkv_weight, get_dflash_attention_sliding_window_size, get_dflash_layer_types, @@ -219,6 +223,19 @@ def kv_proj_only( k, v = kv.split([self.kv_size, self.kv_size], dim=-1) return k, v + can_slice_aiter_qkv_weight, _ = ( + can_dflash_slice_aiter_trans_qkv_weight(self.qkv_proj) + ) + if can_slice_aiter_qkv_weight: + kv_slice = slice(self.q_size, self.q_size + 2 * self.kv_size) + weight = self.qkv_proj.weight[:, kv_slice] + bias = ( + self.qkv_proj.bias[kv_slice] if self.qkv_proj.bias is not None else None + ) + kv = rocm_aiter_swizzle_hipb_unquantized_gemm(hidden_states, weight, bias) + k, v = kv.split([self.kv_size, self.kv_size], dim=-1) + return k, v + # Fallback: compute full QKV and discard Q (keeps compatibility with quantized weights). qkv, _ = self.qkv_proj(hidden_states) _, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index dd1c36ad8eb9..c2ddc86f90a3 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -524,15 +524,16 @@ def parse_dflash_draft_config(*, draft_hf_config: Any) -> DFlashDraftConfig: def can_dflash_slice_qkv_weight(qkv_proj: Any) -> Tuple[bool, str]: """Validate whether DFlash can slice KV weights from a fused QKV linear layer.""" quant_method = getattr(qkv_proj, "quant_method", None) + weight = getattr(qkv_proj, "weight", None) + if not isinstance(quant_method, UnquantizedLinearMethod): return ( False, "quantized qkv_proj is not supported for this path " f"(quant_method={type(quant_method).__name__})", ) - if not hasattr(qkv_proj, "weight"): + if weight is None: return False, "qkv weight tensor is missing" - weight = qkv_proj.weight if getattr(weight, "aiter_trans_weight", False): return False, "AITER transposed qkv weight layout is not sliceable" if getattr(weight, "ndim", None) != 2: @@ -551,6 +552,38 @@ def can_dflash_slice_qkv_weight(qkv_proj: Any) -> Tuple[bool, str]: return True, "" +def can_dflash_slice_aiter_trans_qkv_weight(qkv_proj: Any) -> Tuple[bool, str]: + """Validate whether DFlash can column-slice KV from AITER transposed QKV weight.""" + quant_method = getattr(qkv_proj, "quant_method", None) + weight = getattr(qkv_proj, "weight", None) + + if not isinstance(quant_method, UnquantizedLinearMethod): + return ( + False, + "quantized qkv_proj is not supported for AITER column-slice path " + f"(quant_method={type(quant_method).__name__})", + ) + if weight is None: + return False, "qkv weight tensor is missing" + if not getattr(weight, "aiter_trans_weight", False): + return False, "qkv weight is not in AITER transposed layout" + if getattr(weight, "ndim", None) != 2: + return False, f"qkv weight must be 2D, got shape={tuple(weight.shape)}" + + expected_shape = ( + int(getattr(qkv_proj, "input_size")), + int(getattr(qkv_proj, "output_size_per_partition")), + ) + if tuple(weight.shape) != expected_shape: + return ( + False, + "AITER qkv weight layout is not column-sliceable: " + f"expected shape={expected_shape}, got shape={tuple(weight.shape)}", + ) + + return True, "" + + def can_dflash_use_fused_qkv_proj(qkv_proj: Any) -> Tuple[bool, str]: """Validate whether a QKV layer is eligible for DFlash fused KV materialization.""" eligible, reason = can_dflash_slice_qkv_weight(qkv_proj)