Fix qkv_proj slice in dflash draft model

python/sglang/srt/models/dflash.py

@@ -22,13 +22,17 @@
     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
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 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)

python/sglang/srt/speculative/dflash_utils.py

@@ -524,37 +524,63 @@ 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:
+        return False, f"qkv weight must be 2D, got shape={tuple(weight.shape)}"
+
+    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 uses AITER shuffled layout; direct tensor slicing would bypass "
-            "the linear method",
+            "qkv weight layout is not sliceable: "
+            f"expected shape={expected_shape}, got shape={tuple(weight.shape)}",
         )
-    if weight.ndim != 2:
+    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_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, "input_size")),
+        int(getattr(qkv_proj, "output_size_per_partition")),
+    )
+    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)})",
+            "AITER qkv weight layout is not column-sliceable: "
+            f"expected shape={expected_shape}, got shape={tuple(weight.shape)}",
         )
+
     return True, ""