From 07a6bb966557dc0e6e5a8efd7511b1e5150f19a8 Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Wed, 29 Apr 2026 23:33:54 -0700 Subject: [PATCH] [None][perf] AutoDeploy: FlashInfer SSM kernel for MTP extend path (trtllm_ssm backend) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the Triton _selective_scan_update_kernel with FlashInfer's selective_state_update_kernel_simple_mtp for the MTP (speculative decoding) extend path in AutoDeploy. Enhance accuracy test with fp8 and nvfp4 variants. Acceptance rate threshold for bf16 tightened: 45% -> 50% (achieves ~53%) Acceptance rate for fp8/nvfp4 set to 40% due to numerical issues (achieves ~44%) The acceptance rate and GSM8K accuracy scores are consistent between Triton and FlashInfer. Triton SSM: 0.329 ms/call × 40 layers = 13.2 ms/iter (decode) FlashInfer: 0.149 ms/call × 40 layers = 5.9 ms/iter (decode) Iter speedup: 30.0 ms → 23.5 ms (1.27x) Throughput: +20% across conc=1–128 Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- .../model_registry/configs/super_v3_mtp.yaml | 8 +- .../mamba/flashinfer_backend_mamba.py | 116 +++++++++++++++--- .../defs/accuracy/references/gsm8k.yaml | 4 + .../defs/accuracy/test_llm_api_autodeploy.py | 68 ++++++++-- .../test_lists/qa/llm_function_core.txt | 6 +- .../test_lists/test-db/l0_dgx_b200.yml | 5 +- .../test_lists/test-db/l0_dgx_h100.yml | 1 + .../mamba/test_flashinfer_mamba_cached_op.py | 1 + 8 files changed, 172 insertions(+), 37 deletions(-) diff --git a/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml b/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml index da839c657a16..e4ea9ab6ee0a 100644 --- a/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml +++ b/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml @@ -9,7 +9,7 @@ attn_backend: trtllm model_factory: AutoModelForCausalLM skip_loading_weights: false cuda_graph_config: - batch_sizes: [1, 2, 4, 8, 16, 24, 32, 64, 128, 256, 320, 384] + batch_sizes: [1, 2, 4, 8, 16, 24, 32, 64, 128] kv_cache_config: # tunable mamba cache dtype # --> use float32 for accuracy and default (auto) for speed @@ -50,10 +50,10 @@ transforms: fuse_mamba_a_log: stage: post_load_fusion enabled: true - # Triton SSM + causal conv are required for MTP as currently they are the only backends - # that support speculative mamba state caching. + # FlashInfer SSM is used for the MTP extend path; Triton causal conv is still required + # for speculative Mamba state caching in this configuration. insert_cached_ssm_attention: - backend: triton_ssm + backend: flashinfer_ssm insert_cached_causal_conv: backend: triton_causal_conv fuse_nvfp4_moe: diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py index 7232afd17c6a..67ad432a2231 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py @@ -16,19 +16,30 @@ from typing import List, Optional import torch +from flashinfer.mamba import selective_state_update as _flashinfer_ssm_update from torch.fx import Node from ..._compat import KvCacheConfig -from ..attention_interface import AttentionRegistry, BatchInfo, MHACallable, ResourceHandlerDict +from ..attention_interface import ( + AttentionRegistry, + BatchInfo, + MHACallable, + ResourceHandlerDict, + SpecSSMResourceHandler, +) from .mamba_backend_common import ( BaseBackendSSM, _flatten_ssm_inputs, _prepare_ssm_decode_inputs, + _prepare_ssm_grouped_state_update_inputs, _run_ssm_prefill, ) -@torch.library.custom_op("auto_deploy::flashinfer_cached_ssm", mutates_args=("ssm_state_cache",)) +@torch.library.custom_op( + "auto_deploy::flashinfer_cached_ssm", + mutates_args=("ssm_state_cache", "intermediate_ssm_state_cache"), +) def _flashinfer_cached_ssm( # INPUTS (dense but may be flattened across sequences) hidden_states: torch.Tensor, # [b, s, num_heads, head_dim] @@ -50,6 +61,9 @@ def _flashinfer_cached_ssm( seq_idx_prefill: torch.Tensor, # [1, num_prefill_tokens] # CACHES ssm_state_cache: torch.Tensor, # [max_batch_size, num_heads, head_dim, ssm_state_size] + intermediate_ssm_state_cache: Optional[ + torch.Tensor + ], # [spec_state_size, max_draft_len+1, num_heads, head_dim, d_state] # CONSTANTS time_step_limit: List[float], chunk_size: int, @@ -60,9 +74,10 @@ def _flashinfer_cached_ssm( ) ssm_state_size = B.shape[3] batch_info = BatchInfo(batch_info_host) - num_prefill, _, num_decode = batch_info.get_num_sequences() - num_prefill_tokens, _, num_decode_tokens = batch_info.get_num_tokens() - num_total_tokens = num_prefill_tokens + num_decode_tokens + num_prefill, num_extend, num_decode = batch_info.get_num_sequences() + num_prefill_tokens, num_extend_tokens, num_decode_tokens = batch_info.get_num_tokens() + num_total_tokens = num_prefill_tokens + num_extend_tokens + num_decode_tokens + if out is not None: preallocated_ssm_out = out.view(bs, num_heads, head_dim) else: @@ -72,6 +87,7 @@ def _flashinfer_cached_ssm( device=hidden_states.device, ) + # PREFILL _run_ssm_prefill( hs_flat, B_flat, @@ -94,6 +110,71 @@ def _flashinfer_cached_ssm( preallocated_ssm_out[:num_prefill_tokens].unsqueeze(0), ) + # EXTEND: multi-token MTP verification path, writes intermediate SSM states + extend_inputs = _prepare_ssm_grouped_state_update_inputs( + hs_flat, + B_flat, + C_flat, + dt_flat, + A, + D, + dt_bias, + slot_idx, + seq_start=num_prefill, + token_start=num_prefill_tokens, + num_seq=num_extend, + num_tokens=num_extend_tokens, + num_heads=num_heads, + head_dim=head_dim, + ssm_state_size=ssm_state_size, + ) + + if extend_inputs is not None: + tokens_per_extend = num_extend_tokens // num_extend + if intermediate_ssm_state_cache.size(1) < tokens_per_extend: + raise RuntimeError( + "flashinfer_cached_ssm: intermediate_ssm_state_cache is too small " + f"for extend branch (size1={intermediate_ssm_state_cache.size(1)}, " + f"tokens_per_extend={tokens_per_extend})" + ) + ( + slot_idx_extend, + x_extend, + B_extend, + C_extend, + dt_extend, + A_full, + D_full, + dt_bias_hp, + ) = extend_inputs + + preallocated_ssm_out_e = preallocated_ssm_out[ + num_prefill_tokens : num_prefill_tokens + num_extend_tokens + ].view(num_extend, tokens_per_extend, num_heads, head_dim) + + intermediate_state_indices = torch.arange( + num_extend, dtype=torch.int32, device=slot_idx_extend.device + ) + _flashinfer_ssm_update( + ssm_state_cache, + x_extend, + dt_extend, + A_full, + B_extend, + C_extend, + D_full, + z=None, + dt_bias=dt_bias_hp, + dt_softplus=True, + state_batch_indices=slot_idx_extend.to(torch.int32), + out=preallocated_ssm_out_e, + disable_state_update=True, + intermediate_states_buffer=intermediate_ssm_state_cache, + cache_steps=tokens_per_extend, + intermediate_state_indices=intermediate_state_indices, + ) + + # DECODE: single-token autoregressive path decode_inputs = _prepare_ssm_decode_inputs( hs_flat, B_flat, @@ -103,8 +184,8 @@ def _flashinfer_cached_ssm( D, dt_bias, slot_idx, - num_prefill, - num_prefill_tokens, + num_prefill + num_extend, + num_prefill_tokens + num_extend_tokens, num_decode, num_decode_tokens, num_heads, @@ -124,28 +205,23 @@ def _flashinfer_cached_ssm( D_full, ) = decode_inputs - import flashinfer - - # FlashInfer needs contiguous x/B/C with 128-byte alignment. - x_decode = x_decode.contiguous() - B_decode = B_decode.contiguous() - C_decode = C_decode.contiguous() - slot_idx_decode_i32 = slot_idx_decode.to(torch.int32) - y_decode = flashinfer.mamba.selective_state_update( + y_decode = _flashinfer_ssm_update( ssm_state_cache, x_decode, dt_hp, A_full, B_decode, C_decode, - D=D_full, + D_full, z=None, dt_bias=dt_bias_hp, dt_softplus=True, state_batch_indices=slot_idx_decode_i32, ) - preallocated_ssm_out[num_prefill_tokens:num_total_tokens].copy_(y_decode) + preallocated_ssm_out[num_prefill_tokens + num_extend_tokens : num_total_tokens].copy_( + y_decode + ) if out is not None: # out is reused across CUDA graph replays with varying num_total_tokens, @@ -179,6 +255,9 @@ def _flashinfer_cached_ssm_fake( seq_idx_prefill: torch.Tensor, # [1, num_prefill_tokens] # CACHES ssm_state_cache: torch.Tensor, # [max_batch_size, num_heads, head_dim, ssm_state_size] + intermediate_ssm_state_cache: Optional[ + torch.Tensor + ], # [spec_state_size, max_draft_len+1, num_heads, head_dim, d_state] # CONSTANTS time_step_limit: List[float], chunk_size: int, @@ -218,4 +297,7 @@ def get_cache_initializers( "Consider using 'triton_ssm' backend instead." ) + ret["intermediate_ssm_state_cache"] = SpecSSMResourceHandler.from_base( + ret["ssm_state_cache"] + ) return ret diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index 13accb492c50..7c88502acb81 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -439,6 +439,10 @@ nvidia/Nemotron-Super-V3: accuracy: 80.85 - spec_dec_algo: MTP accuracy: 92.70 + - quant_algo: FP8 + kv_cache_quant_algo: FP8 + spec_dec_algo: MTP + accuracy: 92.61 - quant_algo: MIXED_PRECISION kv_cache_quant_algo: FP8 spec_dec_algo: MTP diff --git a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py index 790c95b54d8e..7b02ba713f1f 100644 --- a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py +++ b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py @@ -706,26 +706,62 @@ def test_functional_small(self, dtype): @skip_pre_hopper @pytest.mark.parametrize("attn_backend", ["flashinfer", "trtllm"]) @pytest.mark.parametrize( - "world_size", + "model_id, world_size", [ pytest.param( + "bf16", 4, - marks=pytest.mark.skip_less_device_memory(180000), - id="ws4_180gb", + marks=[ + pytest.mark.skip_less_device(4), + pytest.mark.skip_less_device_memory(180000), + ], + id="bf16_ws4_180gb", ), pytest.param( + "fp8", + 4, + marks=[ + pytest.mark.skip_less_device(4), + pytest.mark.skip_less_device_memory(80000), + ], + id="fp8_ws4_80gb", + ), + pytest.param( + "nvfp4", + 4, + marks=[ + pytest.mark.skip_less_device(4), + pytest.mark.skip_less_device_memory(80000), + skip_pre_blackwell, + ], + id="nvfp4_ws4_80gb", + ), + pytest.param( + "fp8", + 8, + marks=[ + pytest.mark.skip_less_device(8), + pytest.mark.skip_less_device_memory(80000), + ], + id="fp8_ws8_80gb", + ), + pytest.param( + "nvfp4", 8, - marks=pytest.mark.skip_less_device_memory(80000), - id="ws8_80gb", + marks=[ + pytest.mark.skip_less_device(8), + pytest.mark.skip_less_device_memory(80000), + skip_pre_blackwell, + ], + id="nvfp4_ws8_80gb", ), ], ) - def test_mtp(self, world_size, attn_backend): - if get_device_count() < world_size: - pytest.skip(f"Not enough devices for world_size={world_size}") + def test_mtp(self, world_size, attn_backend, model_id): - model_path = self.MODEL_PATHS["bf16"] + model_path = self.MODEL_PATHS[model_id] kwargs = {} + # TODO: gate for bf16 only after replay lands low_memory_overrides( kwargs, max_batch_size=8, @@ -740,8 +776,8 @@ def test_mtp(self, world_size, attn_backend): kwargs["compile_backend"] = "torch-simple" print( - f"SuperV3 MTP params: world_size={world_size}, model_path={model_path}" - ) + f"SuperV3 MTP params: model_id={model_id}, world_size={world_size}, " + f"model_path={model_path}") print(f"kwargs: {kwargs}") mtp_yaml = str( @@ -759,9 +795,17 @@ def test_mtp(self, world_size, attn_backend): enable_iter_perf_stats=True, **kwargs, ) as llm: + _set_quant_config(llm, model_id) + if model_id == "nvfp4": + llm.args.quant_config.quant_algo = QuantAlgo.MIXED_PRECISION + print_memory_usage("after engine build") + task = GSM8K(self.MODEL_NAME) task.evaluate(llm) - self.check_acceptance_rate(llm, min_acceptance_rate=0.45) + # bf16 acceptance is stable; fp8/nvfp4 have higher variance due to + # arithmetic rounding, so use a lower threshold for quantized models. + min_rate = 0.50 if model_id == "bf16" else 0.40 + self.check_acceptance_rate(llm, min_acceptance_rate=min_rate) print_memory_usage("after evaluation") diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 6ac7f64897b9..2b4e34264144 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -138,8 +138,10 @@ accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[nvfp4-1- accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[nvfp4-4-attn_dp_on-trtllm] accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_functional_small[bf16] accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_functional_small[fp8] -accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[ws4_180gb-flashinfer] -accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[ws4_180gb-trtllm] +accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[bf16_ws4_180gb-flashinfer] +accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[bf16_ws4_180gb-trtllm] +accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[fp8_ws4_80gb-trtllm] +accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[nvfp4_ws4_80gb-trtllm] accuracy/test_llm_api_autodeploy.py::TestNemotronV2::test_auto_dtype[False] accuracy/test_llm_api_autodeploy.py::TestNemotronV2::test_auto_dtype[True] accuracy/test_llm_api_autodeploy.py::TestNemotronV2::test_fp8[True] diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 3b733f48ddd4..6f39924a0289 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -357,8 +357,9 @@ l0_dgx_b200: - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[fp8-4-attn_dp_off-trtllm] - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[fp8-4-attn_dp_on-trtllm] - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[nvfp4-4-attn_dp_on-trtllm] - - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[ws4_180gb-flashinfer] - - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[ws4_180gb-trtllm] + - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[bf16_ws4_180gb-flashinfer] + - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[bf16_ws4_180gb-trtllm] + - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[nvfp4_ws4_80gb-trtllm] - accuracy/test_llm_api_autodeploy.py::TestNemotronUltraV3::test_accuracy[nvfp4-4] - accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[nvidia_Llama-3.1-8B-Instruct-NVFP4-True] # ------------- AutoDeploy Perf Sanity --------------- diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index b3b4264cf042..746bdbd40a49 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -374,6 +374,7 @@ l0_dgx_h100: - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[bf16-4-attn_dp_on-trtllm] - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[fp8-4-attn_dp_off-trtllm] - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[fp8-4-attn_dp_on-trtllm] + - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[fp8_ws4_80gb-trtllm] - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_attention_dp[4] - accuracy/test_llm_api_autodeploy.py::TestQwen3_5_397B_MoE::test_bf16_small[4] - accuracy/test_llm_api_autodeploy.py::TestGemma4MoE::test_bf16 diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py index 8a29a0167217..82614a5e1395 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py @@ -104,6 +104,7 @@ def test_flashinfer_decode_matches_triton(mamba_env): None, # seq_idx_prefill # CACHES ssm_state_cache_flashinfer, + None, # intermediate_ssm_state_cache (not used in decode-only path) # CONSTANTS time_step_limit, chunk_size,