[Perf] DFLASH: route draft bf16 GEMMs through tuned dispatch on GB300

python/sglang/kernels/ops/gemm/cutedsl_bf16_gemm.py

@@ -1370,6 +1370,9 @@ def _pick_tactic(m: int, n: int, k: int) -> int:
     (8192, 1024): (1, 64),  # TP8 o_proj
     (512, 8192): (256, 512),  # TP8 pooler slice
     (1024, 4096): (1, 512),  # TP8 unpooler slice
+    (14336, 8192): (24, 32),  # hidden 8192 / inter 28672 TP4 gate_up (DFLASH draft)
+    (50512, 8192): (17, 32),  # DFLASH draft TP4 candidate head
+    (7168, 8192): (24, 32),  # hidden 8192 / inter 28672 TP8 gate_up (DFLASH draft)
 }
 
 

python/sglang/srt/layers/quantization/unquant.py

@@ -223,6 +223,8 @@ def is_gemv(self) -> bool:
     (8, 3584, 2560): (64, 8, 2, 6),
     (8, 4096, 2560): (64, 8, 2, 6),
     (8, 4120, 2560): (64, 8, 2, 6),
+    # GB300: DFLASH draft candidate head (202048-vocab org shard), TP8.
+    (7, 25256, 8192): (64, 8, 1, 12),
 }
 
 # FlashInfer warp split-K (output_tile, token_tile, k_tile, stages, b_loader_warps),
@@ -267,6 +269,31 @@ def is_gemv(self) -> bool:
     (2, 1024, 4096): (16, 8, 256, 16, 2),
     (24, 1024, 4096): (16, 16, 256, 10, 2),
     (32, 1024, 4096): (16, 16, 256, 10, 2),
+    # GB300, measured on the DFLASH draft of Muse Spark
+    (8, 2048, 8192): (16, 8, 256, 16, 2),
+    (16, 2048, 8192): (16, 16, 256, 13, 2),
+    (24, 2048, 8192): (16, 32, 256, 9, 2),
+    (32, 2048, 8192): (32, 16, 256, 8, 2),
+    # hidden 8192 / intermediate 28672 / 64q8kv heads (also Llama-3-70B-class) at TP4:
+    # gate_up, down, qkv; the qkv rows replace the split-K tactics above.
+    # Then the TP4 candidate head (202048-vocab org shard).
+    (8, 14336, 8192): (16, 8, 128, 12, 2),
+    (16, 14336, 8192): (32, 16, 256, 9, 2),
+    (8, 8192, 7168): (32, 8, 128, 10, 1),
+    (16, 8192, 7168): (32, 16, 256, 4, 2),
+    (8, 2560, 8192): (32, 8, 256, 11, 1),
+    (16, 2560, 8192): (32, 16, 256, 9, 1),
+    (7, 50512, 8192): (16, 8, 256, 9, 2),
+    (14, 50512, 8192): (16, 16, 256, 4, 2),
+    # Same model shapes at TP8: gate_up, down, qkv.
+    (8, 7168, 8192): (16, 8, 256, 6, 2),
+    (16, 7168, 8192): (16, 16, 256, 4, 2),
+    (8, 8192, 3584): (32, 8, 256, 5, 2),
+    (16, 8192, 3584): (32, 16, 128, 8, 1),
+    (8, 1280, 8192): (16, 8, 256, 11, 2),
+    (16, 1280, 8192): (16, 16, 256, 13, 2),
+    (24, 1280, 8192): (32, 8, 256, 9, 2),
+    (32, 1280, 8192): (32, 16, 256, 9, 2),
 }
 
 

python/sglang/srt/models/dflash.py

@@ -30,6 +30,10 @@
     LogitsProcessorOutput,
     should_apply_lm_head_quant_method,
 )
+from sglang.srt.layers.quantization.unquant import (
+    bf16_gemm_dispatch,
+    get_bf16_gemm_backend,
+)
 from sglang.srt.layers.radix_attention import AttentionType, RadixAttention
 from sglang.srt.layers.rotary_embedding import get_rope
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
@@ -94,7 +98,16 @@ def _project_candidate_logits(
     """Project draft hiddens through the target head, restricted to the org vocab."""
     if not use_quant_head:
         weight = lm_head.weight
-        return torch.matmul(hidden.to(weight.dtype), weight[:num_org].T)
+        hidden = hidden.to(weight.dtype)
+        # Tuned bf16 dispatch (kernel chosen by GEMM shape); bf16 CUDA only.
+        backend = get_bf16_gemm_backend()
+        if (
+            (backend.is_cutedsl() or backend.is_gemv())
+            and hidden.is_cuda
+            and weight.dtype == torch.bfloat16
+        ):
+            return bf16_gemm_dispatch(hidden, weight[:num_org], None)
+        return torch.matmul(hidden, weight[:num_org].T)
     # A packed weight can't be row-sliced to the org vocab like the dense path,
     # and flashinfer's radix top-k rejects the crop view (non-contiguous), so
     # mask the padded tail out of the top-k instead.
@@ -558,7 +571,7 @@ def __init__(
         base_kernel = torch.zeros(2, self.taps, hidden_size)
         base_kernel[:, 0] = 1.0
         self.base_kernel = nn.Parameter(base_kernel)
-        self.kernel_projection = nn.Linear(
+        self.kernel_projection = ReplicatedLinear(
             hidden_size, 2 * self.taps * self.num_groups, bias=False
         )
 
@@ -579,7 +592,7 @@ def _convolve(self, hidden_states, delta, side: int) -> torch.Tensor:
         )
 
     def prepare(self, hidden_states: torch.Tensor):
-        coefficients = self.kernel_projection(hidden_states).reshape(
+        coefficients = self.kernel_projection(hidden_states)[0].reshape(
             *hidden_states.shape[:-1], 2, self.taps, self.num_groups
         )
         return (