[None][feat] Add AutoDeploy support for StepFun Step-3.7-Flash

examples/auto_deploy/cookbooks/step_3.7_flash_trtllm_cookbook.ipynb

@@ -0,0 +1,267 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Deploying Step-3.7-Flash with TensorRT-LLM\n",
+    "\n",
+    "This notebook walks you through deploying the `stepfun-ai/Step-3.7-Flash` model (text generation path) using TensorRT-LLM.\n",
+    "\n",
+    "[TensorRT-LLM](https://nvidia.github.io/TensorRT-LLM/) is NVIDIA's open-source library for accelerating and optimizing LLM inference on NVIDIA GPUs. Support for Step-3.7-Flash is enabled through the AutoDeploy workflow. More details about AutoDeploy can be found [here](https://nvidia.github.io/TensorRT-LLM/torch/auto_deploy/auto-deploy.html).\n",
+    "\n",
+    "**Model Resources:**\n",
+    "- [HuggingFace Model Card](https://huggingface.co/stepfun-ai/Step-3.7-Flash)\n",
+    "- [Technical Blog](https://static.stepfun.com/blog/step-3.7-flash/)\n",
+    "- [StepFun API Platform](https://platform.stepfun.ai)\n",
+    "- [Discord Community](https://discord.gg/RcMJhNVAQc)\n",
+    "\n",
+    "**Model Highlights:**\n",
+    "- 198B-parameter sparse Mixture-of-Experts (MoE) model, ~11B active parameters per token\n",
+    "- 288 routed experts (top-8) plus a shared expert; 45 decoder layers\n",
+    "- Mixed full / sliding-window grouped-query attention with a head-wise attention gate\n",
+    "- 256K token context length\n",
+    "- Three reasoning levels (low / medium / high) and tool-calling support\n",
+    "- Apache 2.0 License\n",
+    "\n",
+    "> **Note:** Step-3.7-Flash is a vision-language model. AutoDeploy onboards the **text generation path**; the vision encoder is not deployed through this workflow.\n",
+    "\n",
+    "**Prerequisites:**\n",
+    "- 8x NVIDIA GPUs with recent drivers (BF16 weights are ~400 GB total) and CUDA 12.x\n",
+    "- Python 3.10+\n",
+    "- TensorRT-LLM ([container](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tensorrt-llm/containers/release) or pip install)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Prerequisites & Environment\n",
+    "\n",
+    "Set up a containerized environment for TensorRT-LLM by running the following command in a terminal:\n",
+    "\n",
+    "```shell\n",
+    "docker run --rm -it --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --gpus=all -p 8000:8000 nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc1\n",
+    "```\n",
+    "\n",
+    "You now have TensorRT-LLM set up!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# If pip not found\n",
+    "!python -m ensurepip --default-pip"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install torch openai"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Verify GPU\n",
+    "\n",
+    "Check that CUDA is available and the GPUs are detected correctly. Step-3.7-Flash requires 8 GPUs for the BF16 weights."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Environment check\n",
+    "import sys\n",
+    "\n",
+    "import torch\n",
+    "\n",
+    "print(f\"Python: {sys.version}\")\n",
+    "print(f\"CUDA available: {torch.cuda.is_available()}\")\n",
+    "print(f\"Num GPUs: {torch.cuda.device_count()}\")\n",
+    "\n",
+    "if torch.cuda.is_available():\n",
+    "    for i in range(torch.cuda.device_count()):\n",
+    "        print(f\"GPU[{i}]: {torch.cuda.get_device_name(i)}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## OpenAI-Compatible Server\n",
+    "\n",
+    "Start a local OpenAI-compatible server with TensorRT-LLM via the terminal, within the running docker container.\n",
+    "\n",
+    "Ensure that the following commands are executed from the docker terminal.\n",
+    "\n",
+    "Start with the Step-3.7-Flash YAML here: `examples/auto_deploy/model_registry/configs/step-3.7-flash.yaml`"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Load the Model\n",
+    "\n",
+    "Launch the TensorRT-LLM server with Step-3.7-Flash:\n",
+    "\n",
+    "```shell\n",
+    "trtllm-serve \"stepfun-ai/Step-3.7-Flash\" \\\n",
+    "  --host 0.0.0.0 \\\n",
+    "  --port 8000 \\\n",
+    "  --backend _autodeploy \\\n",
+    "  --trust_remote_code \\\n",
+    "  --extra_llm_api_options examples/auto_deploy/model_registry/configs/step-3.7-flash.yaml\n",
+    "```\n",
+    "\n",
+    "The same standalone config YAML can be validated locally with `build_and_run_ad.py`:\n",
+    "\n",
+    "```shell\n",
+    "python examples/auto_deploy/build_and_run_ad.py \\\n",
+    "  --model stepfun-ai/Step-3.7-Flash \\\n",
+    "  --args.yaml-extra examples/auto_deploy/model_registry/configs/step-3.7-flash.yaml\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Your server is now running!"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Use the API\n",
+    "\n",
+    "Use the OpenAI-compatible client to send requests to the TensorRT-LLM server."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from openai import OpenAI\n",
+    "\n",
+    "# Setup client\n",
+    "BASE_URL = \"http://0.0.0.0:8000/v1\"\n",
+    "API_KEY = \"null\"\n",
+    "client = OpenAI(base_url=BASE_URL, api_key=API_KEY)\n",
+    "\n",
+    "MODEL_ID = \"stepfun-ai/Step-3.7-Flash\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Basic chat completion\n",
+    "print(\"Chat Completion Example\")\n",
+    "print(\"=\" * 50)\n",
+    "\n",
+    "response = client.chat.completions.create(\n",
+    "    model=MODEL_ID,\n",
+    "    messages=[\n",
+    "        {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
+    "        {\"role\": \"user\", \"content\": \"What is 15% of 85? Show your reasoning.\"},\n",
+    "    ],\n",
+    "    temperature=1.0,\n",
+    "    top_p=0.95,\n",
+    "    max_tokens=512,\n",
+    ")\n",
+    "\n",
+    "print(\"Response:\")\n",
+    "print(response.choices[0].message.content)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Streaming chat completion\n",
+    "print(\"Streaming response:\")\n",
+    "print(\"=\" * 50)\n",
+    "\n",
+    "stream = client.chat.completions.create(\n",
+    "    model=MODEL_ID,\n",
+    "    messages=[\n",
+    "        {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
+    "        {\"role\": \"user\", \"content\": \"What are the first 5 prime numbers?\"},\n",
+    "    ],\n",
+    "    temperature=0.7,\n",
+    "    max_tokens=1024,\n",
+    "    stream=True,\n",
+    ")\n",
+    "\n",
+    "for chunk in stream:\n",
+    "    if chunk.choices[0].delta.content:\n",
+    "        print(chunk.choices[0].delta.content, end=\"\", flush=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Evaluation Parameters\n",
+    "\n",
+    "For optimal results, use the following parameters based on your task:\n",
+    "\n",
+    "**Default Settings (Reasoning Tasks)**\n",
+    "- `temperature`: 1.0\n",
+    "- `top_p`: 0.95\n",
+    "- `max_tokens`: up to 8192 (model supports a 256K context)\n",
+    "\n",
+    "**Deterministic Tasks**\n",
+    "- `temperature`: 0\n",
+    "- `max_tokens`: 16384\n",
+    "\n",
+    "Step-3.7-Flash exposes three reasoning levels (low / medium / high); see the model card for how to select them via the chat template."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Additional Resources\n",
+    "\n",
+    "- [TensorRT-LLM Documentation](https://nvidia.github.io/TensorRT-LLM/)\n",
+    "- [AutoDeploy Guide](https://nvidia.github.io/TensorRT-LLM/torch/auto_deploy/auto-deploy.html)\n",
+    "- [Step-3.7-Flash on HuggingFace](https://huggingface.co/stepfun-ai/Step-3.7-Flash)\n",
+    "- [Step-3.7-Flash Technical Blog](https://static.stepfun.com/blog/step-3.7-flash/)\n",
+    "- [StepFun Discord Community](https://discord.gg/RcMJhNVAQc)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

examples/auto_deploy/model_registry/configs/step-3.7-flash.yaml

@@ -0,0 +1,42 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Self-contained AutoDeploy config for StepFun Step-3.7-Flash (text decoder).
+#
+# Step-3.7-Flash is a vision-language model; AutoDeploy onboards the text generation
+# path (model_type step3p7 / text config step3p5). The custom modeling code lives in
+# tensorrt_llm/_torch/auto_deploy/models/custom/modeling_step3p7.py and is selected by
+# registering Step3p7ForCausalLM under the Step3p7Config config class.
+#
+# 45 decoder layers (3 dense + 42 MoE, 288 routed experts top-8 + shared expert),
+# GQA with mixed full/sliding attention and a head-wise attention gate. BF16 weights.
+runtime: trtllm
+compile_backend: torch-cudagraph
+model_factory: AutoModelForCausalLM
+attn_backend: flashinfer
+max_seq_len: 4096
+max_num_tokens: 4096
+max_batch_size: 64
+world_size: 8
+enable_chunked_prefill: true
+cuda_graph_config:
+  batch_sizes: [1, 2, 4, 8, 16, 32, 64]
+kv_cache_config:
+  dtype: bfloat16
+  enable_block_reuse: false
+  free_gpu_memory_fraction: 0.7
+  tokens_per_block: 64
+model_kwargs:
+  torch_dtype: bfloat16
+# Hint-driven sharding IR: the modeling code carries explicit sharding hints
+# (tp_mode / layer_type / tp_scaled_dim / all_reduce), so disable the legacy
+# heuristic sharding and apply the hints instead.
+transforms:
+  detect_sharding:
+    enabled: false
+  sharding_transform_executor:
+    enabled: false
+  apply_sharding_hints:
+    enabled: true
+  gather_logits_before_lm_head:
+    enabled: true

examples/auto_deploy/model_registry/models.yaml

@@ -404,6 +404,18 @@ models:
 # - name: XiaomiMiMo/MiMo-V2-Flash
 #   config_id: default_ws_8
 #   yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml']
+# --- StepFun Step-3.7-Flash (Feb 2026) ---
+# Text decoder of the Step-3.7-Flash VLM (model_type step3p7 / step3p5). 45 layers,
+# 288 routed experts top-8 + shared expert, mixed full/sliding GQA, head-wise attn gate.
+- name: stepfun-ai/Step-3.7-Flash
+  config_id: step_3_7_flash
+  yaml_extra: ['step-3.7-flash.yaml']
+- name: stepfun-ai/Step-3.7-Flash-FP8
+  config_id: step_3_7_flash
+  yaml_extra: ['step-3.7-flash.yaml']
+- name: stepfun-ai/Step-3.7-Flash-NVFP4
+  config_id: step_3_7_flash
+  yaml_extra: ['step-3.7-flash.yaml']
 # --- Kimi-K2.5 (Jan 2026) ---
 # TypeError: 'NoneType' object is not subscriptable
 # - name: moonshotai/Kimi-K2.5

tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py

@@ -416,6 +416,11 @@ def flashinfer_mha_with_cache(
             kv_layout=_GlobalFlashInferPlanner.kv_layout,
         )
 
+    if num_prefill > 0 and not read_cache_only:
+        # FlashInfer planning depends on the preceding paged-KV append completing.
+        # Keep the same append-before-plan synchronization used by the PyTorch backend.
+        torch.cuda.current_stream().synchronize()
+
     bs = b * s
     if out is not None:
         y = out.view(-1, n_heads, head_dim)

tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py

@@ -63,6 +63,7 @@
     "modeling_skywork_r1v2": ["SkyworkR1V2ForConditionalGeneration"],
     "modeling_smollm3": ["SmolLM3ForCausalLM"],
     "modeling_starcoder2": ["Starcoder2ForCausalLM"],
+    "modeling_step3p7": ["Step3p7ForCausalLM"],
 }
 
 # AD_USE_IR_MODELS: opt-in flag for staging additional ``_ir.py`` modeling