From 0c29c8fecee11f6e9cbe6a69142ceeb44fc41645 Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Thu, 23 Jul 2026 05:21:59 +0800 Subject: [PATCH] Bump FlashInfer to 0.6.15.post1 (#31927) Co-authored-by: Brayden Zhong --- docker/Dockerfile | 2 +- python/pyproject.toml | 2 +- python/sglang/srt/entrypoints/engine.py | 2 +- .../srt/layers/attention/dsa_backend.py | 22 ++++++++++ .../layers/attention/trtllm_mla_backend.py | 41 +++++++++++++++++++ .../moe/moe_runner/flashinfer_cutedsl.py | 17 +++++++- python/sglang/srt/models/deepseek_v2.py | 8 ---- python/sglang/srt/utils/common.py | 2 +- .../attention_methods/dsa_attention.py | 1 + .../attention_methods/mla_attention.py | 1 + .../models_e2e/test_dsa_glm52_nvfp4_tp_mtp.py | 2 +- 11 files changed, 85 insertions(+), 15 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index f318a292c..6001be1f6 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -19,7 +19,7 @@ ARG PIP_DEFAULT_INDEX ARG UBUNTU_MIRROR ARG GITHUB_ARTIFACTORY=github.com ARG INSTALL_FLASHINFER_JIT_CACHE=0 -ARG FLASHINFER_VERSION=0.6.14 +ARG FLASHINFER_VERSION=0.6.15.post1 ARG MOONCAKE_VERSION=0.3.11.post1 ARG MSCCLPP_VERSION=sglang-v0.9.1 #if need other arg please add in MOONCAKE_COMPILE_ARG diff --git a/python/pyproject.toml b/python/pyproject.toml index f0146ceba..4131966a2 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -31,7 +31,7 @@ dependencies = [ "einops", "fastapi", "flash-attn-4>=4.0.0b18", - "flashinfer_python[cu13]==0.6.14", # keep it aligned with jit-cache version in Dockerfile + "flashinfer_python[cu13]==0.6.15.post1", # keep it aligned with jit-cache version in Dockerfile "gguf", "helion==0.2.6", "humming-kernels[cu13]==0.1.10", diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 5c158535e..888542300 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -1304,7 +1304,7 @@ def _set_envs_and_config(server_args: ServerArgs): if server_args.attention_backend == "flashinfer": assert_pkg_version( "flashinfer_python", - "0.6.14", + "0.6.15.post1", "Please uninstall the old version and " "reinstall the latest version by following the instructions " "at https://docs.flashinfer.ai/installation.html.", diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index d8883d29b..f8871b74c 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -56,6 +56,10 @@ from sglang.srt.layers.attention.dsa.utils import ( pad_dsa_cache_seqlens, should_use_dsa_fused_topk, ) +from sglang.srt.layers.attention.trtllm_mla_backend import ( + grow_multi_ctas_kv_counter_buffer_if_needed, + make_persistent_multi_ctas_kv_counter_buffer, +) from sglang.srt.layers.utils.cp_utils import ( cp_all_gather_rerange_output, cp_split_and_rebuild_position, @@ -496,8 +500,16 @@ class DeepseekSparseAttnBackend( device=model_runner.device, ), ) + self._multi_ctas_kv_counter_buffer = ( + make_persistent_multi_ctas_kv_counter_buffer( + torch.device(self.device), + self.num_q_heads, + max_batch_size=model_runner.max_running_requests, + ) + ) else: self.workspace_buffer = None + self._multi_ctas_kv_counter_buffer = None def _make_aiter_dsa_decode_metadata_buffer( self, @@ -2917,6 +2929,15 @@ class DeepseekSparseAttnBackend( batch_size = page_table_1.shape[0] _, num_heads, head_dim = q_all.shape + self._multi_ctas_kv_counter_buffer = ( + grow_multi_ctas_kv_counter_buffer_if_needed( + self._multi_ctas_kv_counter_buffer, + torch.device(self.device), + self.num_q_heads, + batch_size, + ) + ) + q = q_all.view(batch_size, 1, num_heads, head_dim) kv = kv_cache.view(-1, 1, self.real_page_size, self.kv_cache_dim) block_tables = page_table_1.unsqueeze(1) @@ -2945,6 +2966,7 @@ class DeepseekSparseAttnBackend( bmm1_scale=bmm1_scale, backend="trtllm-gen", skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(), + multi_ctas_kv_counter_buffer=self._multi_ctas_kv_counter_buffer, ) return out diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index d69c6c99a..ebd32a361 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -61,6 +61,35 @@ DEFAULT_WORKSPACE_SIZE_MB = 150 # Memory workspace size in MB # compute the LCM with other padding constraints. TRTLLM_BLOCK_CONSTRAINT = 128 +TRTLLM_MLA_MAX_BATCH_SIZE = 8192 + + +def _multi_ctas_kv_counter_bytes( + device: torch.device, num_q_heads: int, batch_size: int +) -> int: + sm_count = flashinfer.utils.get_device_sm_count(device) + return flashinfer.utils.get_trtllm_gen_multi_ctas_kv_counter_bytes( + batch_size, num_q_heads, sm_count + ) + + +def make_persistent_multi_ctas_kv_counter_buffer( + device: torch.device, num_q_heads: int, max_batch_size: int +) -> torch.Tensor: + num_bytes = _multi_ctas_kv_counter_bytes( + device, num_q_heads, max(TRTLLM_MLA_MAX_BATCH_SIZE, max_batch_size) + ) + return torch.zeros(num_bytes, dtype=torch.uint8, device=device) + + +def grow_multi_ctas_kv_counter_buffer_if_needed( + buffer: torch.Tensor, device: torch.device, num_q_heads: int, batch_size: int +) -> torch.Tensor: + required_bytes = _multi_ctas_kv_counter_bytes(device, num_q_heads, batch_size) + if buffer.numel() >= required_bytes: + return buffer + return torch.zeros(required_bytes, dtype=torch.uint8, device=device) + def _quantize_fp8_qkv(q, k, v, layer): q = q.to(torch.float8_e4m3fn) @@ -189,6 +218,14 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): ), ) + self._multi_ctas_kv_counter_buffer = ( + make_persistent_multi_ctas_kv_counter_buffer( + torch.device(self.device), + self.num_q_heads, + max_batch_size=model_runner.max_running_requests, + ) + ) + # CUDA graph state self.decode_cuda_graph_metadata = {} self.decode_cuda_graph_kv_indices = None @@ -642,6 +679,10 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): seq_lens if seq_lens.dtype == torch.int32 else seq_lens.to(torch.int32) ) extra_kwargs = {"backend": self.backend} if self.backend != "trtllm-gen" else {} + if self.backend == "trtllm-gen": + extra_kwargs["multi_ctas_kv_counter_buffer"] = ( + self._multi_ctas_kv_counter_buffer + ) return flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla( query=query, kv_cache=kv_cache, diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py index daa6ba2f9..02366f886 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py @@ -222,6 +222,17 @@ def resolve_cutedsl_standard_scales( return w1_alpha, fc2_input_scale, w2_alpha, used_input_scale +def _cutedsl_wrapper_activation_type(activation: str, activation_type_cls: Any) -> Any: + if activation == "silu": + return activation_type_cls.Swiglu + if activation == "relu2": + return activation_type_cls.Relu2 + raise ValueError( + f"CuteDSL MoE wrapper supports 'silu' (gated) or 'relu2' (non-gated) " + f"activation, got {activation!r}." + ) + + def ensure_cutedsl_wrapper(layer: torch.nn.Module) -> None: """Lazily create CuteDslMoEWrapper and resolve scales on first forward. @@ -237,7 +248,7 @@ def ensure_cutedsl_wrapper(layer: torch.nn.Module) -> None: return try: - from flashinfer import CuteDslMoEWrapper + from flashinfer import ActivationType, CuteDslMoEWrapper except ImportError as e: raise ImportError( "flashinfer_cutedsl backend requires FlashInfer with CuteDSL support. " @@ -284,7 +295,9 @@ def ensure_cutedsl_wrapper(layer: torch.nn.Module) -> None: local_expert_offset=layer.moe_ep_rank * layer.num_local_experts, output_dtype=layer.moe_runner_config.params_dtype, device=str(layer.w13_weight.device), - activation=layer.moe_runner_config.activation, + activation_type=_cutedsl_wrapper_activation_type( + layer.moe_runner_config.activation, ActivationType + ), ) w1_alpha, fc2_input_scale, w2_alpha, used_input_scale = ( diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 0c2745da8..46512b69c 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -461,14 +461,6 @@ class MoEGate(nn.Module): "quark", ): correction_bias_dtype = torch.bfloat16 - # NOTE(kpham-sgl): flashinfer trtllm routing requires a bf16 - # routing_bias; an fp32 bias yields NaN routing on exact ties. - # Mirror the fp8 path's cast. - if ( - quant_config.get_name() == "modelopt_fp4" - and get_moe_runner_backend().is_flashinfer_trtllm() - ): - correction_bias_dtype = torch.bfloat16 self.e_score_correction_bias = nn.Parameter( torch.empty((config.n_routed_experts), dtype=correction_bias_dtype) ) diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 103a89c98..5a1259cb0 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -1911,7 +1911,7 @@ def check_pkg_version_at_least(pkg: str, min_version: str) -> bool: Args: pkg: Package name (distribution name, e.g., "flashinfer-python") - min_version: Minimum version required (e.g., "0.6.14") + min_version: Minimum version required (e.g., "0.6.15.post1") Returns: True if package is installed and version >= min_version, False otherwise diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py index bbb1c53ac..62b6042a0 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py @@ -354,6 +354,7 @@ class DSAMockModelRunner(ModelRunner): triton_attention_split_tile_size=None, ) self.server_args = self._server_args_override.install() + self.max_running_requests = pool_batch_size self.req_to_token_pool = ReqToTokenPool( size=pool_batch_size, max_context_len=max_context_len, diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py index 670e72d37..cfe81c0b0 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py @@ -282,6 +282,7 @@ class MockMLAModelRunner(ModelRunner): triton_attention_split_tile_size=None, ) self.server_args = self._server_args_override.install() + self.max_running_requests = pool_batch_size self.req_to_token_pool = ReqToTokenPool( size=pool_batch_size, max_context_len=max_context_len, diff --git a/test/registered/models_e2e/test_dsa_glm52_nvfp4_tp_mtp.py b/test/registered/models_e2e/test_dsa_glm52_nvfp4_tp_mtp.py index 210112a5a..8f1c3d84d 100644 --- a/test/registered/models_e2e/test_dsa_glm52_nvfp4_tp_mtp.py +++ b/test/registered/models_e2e/test_dsa_glm52_nvfp4_tp_mtp.py @@ -21,7 +21,7 @@ class TestGLM52NVFP4TPMTP( model = "nvidia/GLM-5.2-NVFP4" tp_size = 4 mem_fraction_static = 0.8 - bs_1_speed_thres = 250 + bs_1_speed_thres = 280 extra_server_args = [ "--moe-runner-backend", "flashinfer_trtllm",