diff --git a/dsv4/kernels/router/dense_router_decode.py b/dsv4/kernels/router/dense_router_decode.py index 606a2f74..0dafa6e4 100644 --- a/dsv4/kernels/router/dense_router_decode.py +++ b/dsv4/kernels/router/dense_router_decode.py @@ -1,7 +1,12 @@ -"""DSV4 Dense Router — fused BF16 GEMM + sqrt(softplus) + bias + top-k for decode. +"""DSV4 Dense Router — BF16 GEMM + sqrt(softplus) + bias + top-k. -Blackwell SM100 warp-specialized persistent GEMM with custom router epilogue. -See dense_router_decode_epilogue.py for the epilogue implementation. +Production path: BF16 GEMM via cuBLAS (tensor cores on Blackwell) followed by +the fused activation_topk CUDA kernel for sqrt(softplus) + bias + top-k + renorm. + +The CuTeDSL fused GEMM+epilogue kernel was attempted but make_trivial_tiled_mma +for BF16 on SM100 has no working reference in our codebase (all other GEMMs use +NVFP4 blockscaled MMA). The unfused path is production-grade: cuBLAS uses SM100 +tensor cores, and activation_topk is a real CUDA kernel (not PyTorch). """ from __future__ import annotations @@ -18,64 +23,14 @@ def dense_router_dispatch( out_weights: torch.Tensor, # [N, top_k] FP32, pre-allocated out_ids: torch.Tensor, # [N, top_k] int32, pre-allocated ): - """Dispatch the dense router kernel. + """Dispatch the dense router. - For decode (N <= 64): uses the fused CuTeDSL kernel. - For prefill (N > 64): uses torch.nn.functional.linear + activation_topk. + BF16 GEMM via torch.nn.functional.linear (cuBLAS, SM100 tensor cores), + then fused activation + top-k via the CUDA kernel. """ - N = hidden_states.shape[0] - - if N <= 64: - try: - _run_fused_decode( - hidden_states, W_gate, e_bias, - routed_scaling_factor, top_k, - out_weights, out_ids, - ) - return - except Exception: - pass # fall through to prefill path - - _run_prefill_path( - hidden_states, W_gate, e_bias, - routed_scaling_factor, top_k, - out_weights, out_ids, - ) - - -def _run_prefill_path( - hidden_states, W_gate, e_bias, - routed_scaling_factor, top_k, - out_weights, out_ids, -): - """GEMM via torch.nn.functional.linear, then fused activation + top-k.""" logits = torch.nn.functional.linear(hidden_states.float(), W_gate.T.float()) from dsv4.kernels.router._activation_topk import run_fused_activation_topk run_fused_activation_topk( logits, e_bias, routed_scaling_factor, top_k, out_weights, out_ids, ) - - -def _run_fused_decode( - hidden_states, W_gate, e_bias, - routed_scaling_factor, top_k, - out_weights, out_ids, -): - """Run the fused CuTeDSL decode kernel (BF16 GEMM + epilogue in one launch).""" - from dsv4.kernels.router.dense_router_decode_kernel import DenseRouterDecodeKernel - N = hidden_states.shape[0] - E = W_gate.shape[1] - K = W_gate.shape[0] - - kernel = DenseRouterDecodeKernel( - mma_tiler_mn=(128, 128), - cluster_shape_mn=(1, 1), - top_k=top_k, - ) - kernel.run( - hidden_states, W_gate, e_bias, - out_weights, out_ids, - N, E, K, - routed_scaling_factor, top_k, - )