[None][feat] Support Step-3.7-Flash model#14711
Conversation
📝 WalkthroughWalkthroughAdds complete Step3p7 model support to TensorRT-LLM PyTorch backend: text-only causal LM with Gemma-style attention and routed MoE, multimodal vision tower wrapper, speculative decoding MTP layers, weight preprocessing utilities, and integration tests covering accuracy, API, parity, and reference generation. ChangesStep3p7 Model Integration
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (4)
tests/integration/defs/accuracy/test_step3p7_gsm8k_canary.py (1)
401-473: Test coverage status: sufficient for canary matrix execution, but parser edge cases need follow-up.Coverage is sufficient here for the acceptance-canary execution path (
cg/offandcg/onvia subprocess matrix).
Follow-up outside this PR: add targeted parser assertions for_extract_intedge cases intests/integration/defs/accuracy/test_step3p7_gsm8k_canary.py(or a dedicated adjacent test file) so scoring logic changes are validated independently of model output variance.As per coding guidelines
tests/**: "Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM. Keep feedback actionable: suggest concrete list file names and whether coverage is sufficient, insufficient, or needs follow-up outside the PR."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/integration/defs/accuracy/test_step3p7_gsm8k_canary.py` around lines 401 - 473, Add targeted unit tests covering edge cases for the parser function _extract_int by creating a new test (e.g., test_extract_int_edge_cases) in tests/integration/defs/accuracy/test_step3p7_gsm8k_canary.py or an adjacent test file; call _extract_int with inputs such as None, empty string, strings with leading/trailing whitespace, malformed numeric strings (letters, symbols), signed numbers, decimal-containing strings, extremely large numbers, and negative values, and assert expected return values or raised errors so scoring logic that depends on parsed integers is independently validated from model output variance.tests/integration/defs/models/test_step3p7_llm_api.py (1)
59-97: ⚡ Quick winUse explicit
| Nonefor optional parameters.PEP 484 prohibits implicit
Optional. Per coding guidelines, use| Nonesyntax instead of relying on= Noneto imply optionality.♻️ Suggested fix
-def _emit_pending(reason: str, args: argparse.Namespace, extra: dict = None) -> int: +def _emit_pending(reason: str, args: argparse.Namespace, extra: dict | None = None) -> int:def _emit_blocker( args: argparse.Namespace, blocker_id: str, detail: str, - exc: BaseException = None, - config: str = None, + exc: BaseException | None = None, + config: str | None = None, ) -> int:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/integration/defs/models/test_step3p7_llm_api.py` around lines 59 - 97, The optional parameters should use explicit union-with-None annotations instead of implicit Optional via default None: update _emit_pending's signature to declare extra as dict | None, and update _emit_blocker's signature to declare exc as BaseException | None and config as str | None (keep other types and return types the same and only change the type annotations on the named parameters in _emit_pending and _emit_blocker).tests/integration/defs/models/test_step3p7_parity.py (1)
260-267: ⚡ Quick winUse explicit
| Nonefor optional parameter.Same issue as in
test_step3p7_llm_api.py- PEP 484 prohibits implicitOptional.♻️ Suggested fix
-def _emit_blocker(args, blocker_id: str, detail: str, exc: BaseException = None, **fields) -> int: +def _emit_blocker(args, blocker_id: str, detail: str, exc: BaseException | None = None, **fields) -> int:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/integration/defs/models/test_step3p7_parity.py` around lines 260 - 267, The function _emit_blocker has an optional parameter typed as "exc: BaseException = None" which implicitly suggests Optional; update the signature to use explicit union syntax: change the parameter to "exc: BaseException | None = None" (keeping the default None) so the type clearly expresses optionality; ensure any other references or annotations expecting BaseException are compatible with the new "BaseException | None" type.tests/integration/defs/models/step3p7_hf_reference.py (1)
305-324: 💤 Low valueUnused loop variable
basecan be replaced with_.The
basevariable from the dict iteration is never used inside the loop body.♻️ Suggested fix
- for base, paths in grouped.items(): + for _, paths in grouped.items():🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/integration/defs/models/step3p7_hf_reference.py` around lines 305 - 324, The loop variable "base" in the comprehension for "for base, paths in grouped.items()" is unused; rename it to "_" to indicate it's intentionally ignored (i.e., change the loop header to "for _, paths in grouped.items()") so linters don't flag an unused variable while leaving _shard_grouped_keys, grouped, paths, and the loop body (_remap_key, _matches_layer_target, _maybe_dequant_to_bf16, _materialize_module_param) unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 851-856: The presize helper only runs when eager attention
metadata exists, leaving CUDA-graph capture vulnerable for paths that skip
_general_warmup()/_run_autotuner_warmup()
(draft-model/HELIX/guided-decoder/MambaHybridCacheManager); ensure attention
workspace is always allocated before any CUDA-graph capture by
moving/duplicating the initialization into the capture path: update
_run_cuda_graph_warmup (and other capture entry points) to call a new ensure
routine (or extend _presize_attention_cuda_graph_workspace) that will
create/initialize attn_metadata and allocate/grow attn_metadata.workspace to the
warmed-up size if absent, rather than returning early; reference
_presize_attention_cuda_graph_workspace, _run_cuda_graph_warmup, warmup(),
_general_warmup(), _run_autotuner_warmup(), and attn_metadata.workspace when
making this change.
In `@tests/integration/defs/accuracy/test_step3p7_gsm8k_canary.py`:
- Around line 112-118: Replace broad exception handling with targeted exception
types and proper logging: in the except block that currently does "except
Exception as e:" before returning _FALLBACK_GSM8K_PROBES, change to catch
specific errors expected from the GSM8K loader (e.g., FileNotFoundError,
ValueError, RuntimeError, OSError) and log the exception via the module logger
(logger = logging.getLogger(__name__); logger.warning(..., exc_info=e)) instead
of print; for the locations flagged (the BaseException/bare except patterns at
the ranges referenced) ensure you do not catch BaseException (do not swallow
KeyboardInterrupt/SystemExit) — either re-raise SystemExit/KeyboardInterrupt if
detected or limit to specific exceptions, and avoid silent "except: pass" by
logging the error or handling the specific expected exception.
- Around line 219-225: The integer-extraction fallback currently returns the
first matched token (nums[0]) from the variable cut which often picks up
intermediate chain-of-thought numbers; change the logic to prefer the final
reported integer by using the last match (nums[-1]) before parsing, i.e. update
the try/except block that calls int(nums[0].replace(",", "")) to use
nums[-1].replace(",", "") so the function returns the final answer token instead
of an earlier intermediate value.
- Around line 295-320: The loop uses zip(probes, outputs) which silently
truncates when outputs is shorter than probes; before the loop in the scoring
function, check if len(outputs) != len(probes) and handle it explicitly (e.g.,
log/processLogger or add an error to info and return a non-zero status) so
mismatches are surfaced; update the block that builds info/per_prompt (variables
probes, outputs, per_prompt, info) to abort or mark failure when lengths differ
rather than relying on zip to hide missing outputs.
In `@tests/integration/defs/models/data/step3p7_gsm8k_prompts.json`:
- Line 48: Update the prompt string in the JSON entry where the "prompt" key
contains "How load does it take to download the file?" to correct the typo to
"How long does it take to download the file?"; locate the JSON object with the
"prompt" field in this file and replace only that substring so the rest of the
prompt remains unchanged.
In `@tests/unittest/_torch/modeling/test_modeling_step3p7.py`:
- Around line 23-60: Replace hard failures on missing local checkpoints with
skip-gating: make CHECKPOINT_ROOT configurable from an environment variable
(falling back to the current Path) and use pytest.skip(...) when checkpoint
files are absent rather than pytest.fail; update the checkpoint_path fixture
(and the similar logic referenced around symbols CHECKPOINT_PATHS,
CHECKPOINT_IDS, ALL_CHECKPOINTS_AVAILABLE, SKIP_REASON_NO_CKPT and the other
fixture at lines ~285-330) to call pytest.skip with SKIP_REASON_NO_CKPT if
request.param.exists() is false (or if ALL_CHECKPOINTS_AVAILABLE is false) so
tests are skipped cleanly in environments without the artifacts.
---
Nitpick comments:
In `@tests/integration/defs/accuracy/test_step3p7_gsm8k_canary.py`:
- Around line 401-473: Add targeted unit tests covering edge cases for the
parser function _extract_int by creating a new test (e.g.,
test_extract_int_edge_cases) in
tests/integration/defs/accuracy/test_step3p7_gsm8k_canary.py or an adjacent test
file; call _extract_int with inputs such as None, empty string, strings with
leading/trailing whitespace, malformed numeric strings (letters, symbols),
signed numbers, decimal-containing strings, extremely large numbers, and
negative values, and assert expected return values or raised errors so scoring
logic that depends on parsed integers is independently validated from model
output variance.
In `@tests/integration/defs/models/step3p7_hf_reference.py`:
- Around line 305-324: The loop variable "base" in the comprehension for "for
base, paths in grouped.items()" is unused; rename it to "_" to indicate it's
intentionally ignored (i.e., change the loop header to "for _, paths in
grouped.items()") so linters don't flag an unused variable while leaving
_shard_grouped_keys, grouped, paths, and the loop body (_remap_key,
_matches_layer_target, _maybe_dequant_to_bf16, _materialize_module_param)
unchanged.
In `@tests/integration/defs/models/test_step3p7_llm_api.py`:
- Around line 59-97: The optional parameters should use explicit union-with-None
annotations instead of implicit Optional via default None: update
_emit_pending's signature to declare extra as dict | None, and update
_emit_blocker's signature to declare exc as BaseException | None and config as
str | None (keep other types and return types the same and only change the type
annotations on the named parameters in _emit_pending and _emit_blocker).
In `@tests/integration/defs/models/test_step3p7_parity.py`:
- Around line 260-267: The function _emit_blocker has an optional parameter
typed as "exc: BaseException = None" which implicitly suggests Optional; update
the signature to use explicit union syntax: change the parameter to "exc:
BaseException | None = None" (keeping the default None) so the type clearly
expresses optionality; ensure any other references or annotations expecting
BaseException are compatible with the new "BaseException | None" type.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 361895f8-910a-4093-833b-99b9ea474ba6
📒 Files selected for processing (19)
docs/source/features/speculative-decoding.mddocs/source/models/supported-models.mdtensorrt_llm/_torch/models/__init__.pytensorrt_llm/_torch/models/modeling_speculative.pytensorrt_llm/_torch/models/modeling_step3p7.pytensorrt_llm/_torch/models/modeling_step3p7vl.pytensorrt_llm/_torch/modules/fused_moe/moe_op_backend.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/model_loader.pytensorrt_llm/evaluate/lm_eval.pytensorrt_llm/llmapi/llm_args.pytests/integration/defs/accuracy/test_step3p7_gsm8k_canary.pytests/integration/defs/models/data/step3p7_gsm8k_prompts.jsontests/integration/defs/models/data/step3p7_parity_prompts.jsontests/integration/defs/models/step3p7_hf_reference.pytests/integration/defs/models/test_step3p7_llm_api.pytests/integration/defs/models/test_step3p7_parity.pytests/unittest/_torch/modeling/test_modeling_step3p7.pytests/unittest/llmapi/test_llm_args.py
kaiyux
changed the title
kaiyux
changed the title
yechank-nvidia
left a comment
yechank-nvidia
left a comment
There was a problem hiding this comment.
This PR adds and documents image support for Step3p7ForConditionalGeneration, but the test coverage is still text/config/weight-accounting only.
Do you have plan to add multimodal-related tests?
Also, multimodal hashing looks unsupported/broken for Step3p7 right now. Step3p7VLInputProcessor does not implement get_num_tokens_per_image(), get_mm_token_ids(), or get_mm_special_token_ids(), which are needed by the generic mm hashing path. Since Step3 image spans include structural tokens such as <patch_start>, <patch_end>, <patch_newline>, <im_start>, and <im_end>, the processor should either implement these hooks correctly.
|
@yechank-nvidia Thanks for help review.
Yes will do.
Let me take a look at that, however if that's the case, is there a way to add some kind of checkers? Otherwise I guess developers like me who are not aware of that context will miss, but the model can still run. |
|
Hi @kaiyux, we have multimodal claude skills.md you can take a look at. Also, for the test-related, we have MMMU test(e.g. EXAONE-4.5) and with MMMU, you can explicitly set with chunked_prefill. Also, we have the skeleton test_modeling_multimodal.py which you can do logits matching with HF or sanity check on cuda_graph, kv_cache_reuse, and then multiple modalities. |
@yechank-nvidia Thanks for the information, I'll take a look. By saying add some kind of checkers, I meant if there is a way to raise errors if the required "get_num_tokens_per_image(), get_mm_token_ids(), or get_mm_special_token_ids()" hooks are not implemented, that's another topic though. |
|
/bot run |
|
PR_Github #51508 [ run ] triggered by Bot. Commit: |
|
PR_Github #51508 [ run ] completed with state
|
|
/bot run |
|
PR_Github #51567 [ run ] triggered by Bot. Commit: |
|
PR_Github #51567 [ run ] completed with state
|
|
/bot run |
2 similar comments
|
/bot run |
|
/bot run |
|
PR_Github #51619 [ run ] triggered by Bot. Commit: |
|
PR_Github #51619 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #51709 [ run ] triggered by Bot. Commit: |
|
PR_Github #51709 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #51836 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #51836 [ run ] completed with state
|
|
PR_Github #51862 [ run ] triggered by Bot. Commit: |
|
PR_Github #51862 [ run ] completed with state
|
Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
|
/bot skip --comment "skipping since the only failing test case was moved to qa in the latest commit" |
|
PR_Github #51928 [ skip ] triggered by Bot. Commit: |
|
PR_Github #51928 [ skip ] completed with state |
8 participants
Summary by CodeRabbit
Release Notes
New Features
Bug Fixes
Documentation
Tests
Description
Adds support for StepFun Step-3.7-Flash to the PyTorch backend, covering both the text-only and vision-language variants, plus native MTP speculative decoding and a few supporting fixes.
New model
tensorrt_llm/_torch/models/modeling_step3p7.pyStep3p7ForCausalLM(also registered asStep3p5ForCausalLMfor standalone text use). Hybrid architecture: per-layer attention geometry with partial RoPE, plus a routed MoE block with a shared expert and a sigmoid router with learnable bias (routed_scaling_factorapplied inStep3p7MoE.forward).Step3p7MTP/Step3p7MTPHeadfor speculative decoding, built onSpecDecOneEngineForCausalLM.tensorrt_llm/_torch/models/modeling_step3p7vl.pyStep3p7VLForConditionalGeneration(registered as the HF archStep3p7ForConditionalGeneration) with its ownStep3p7VLInputProcessor. Supports text + image (L + I). The PerceptionEncoder vision tower runs in BF16 even when the text decoder is quantized (FP8 block-scale or NVFP4), and can load from either the FP8 or NVFP4 export. The input processor exposes multimodal-hashing hooks so image embeddings participate in KV-cache reuse.Supporting changes
model_loader.py: whenever aspeculative_configis set, callupdate_spec_config_from_model_config(...)so MTP resolves its layer count and draft length from the loaded checkpoint config before the model runs.speculative/utils.py(update_spec_config_from_model_config): setnum_nextn_predict_layersfrom the model config; when the user leftmax_draft_lenunset (None), resolve it from the checkpoint (vanilla MTP → all MTP heads, MTP-Eagle → 1); when the user set it but it disagrees with the checkpoint for vanilla MTP, clamp to the smaller value and warn; mirror the result intomax_total_draft_tokens.llm_args.py(MTPDecodingConfig): stop defaultingmax_draft_lento1. Leave itNoneas the "use the checkpoint'snum_nextn_predict_layers" sentinel, only validating (> 0) and mirroring tomax_total_draft_tokenswhen the user supplies it — so downstream code can distinguish user-set from defaulted values.modeling_speculative.py: wirestep3p7/step3p5model types intoMTPForCausalLM(→Step3p7MTP).speculative/mtp.py(MTPWorker): clamp generation-requestkv_lens_cudato at leastmtp_num_modulesin place (CUDA-graph safe). The tiny dummy sequences used during generation-step warmup could otherwise drivekv_len < q_len, giving the attention kernel a negative effective KV range and an illegal memory access (seen with Step3p7 MTP + dense sliding-window attention).model_engine.py): add_presize_attention_cuda_graph_workspace(), run after warmup and before capture. The TRTLLM C++ attention op resizes its workspace in place when the requested size grows; inside a CUDA graph capture,resize_()frees the old storage so kernels captured earlier hold a dangling pointer and segfault on replay. Pre-sizingcuda_graph_workspaceto the warmed-up eagerworkspacesize eliminates the in-capture resize.lm_eval.py): add--post_process_fn strip_thinking_mmmu(resolves tostrip_thinking_and_extract_mmmu_answer, which strips<think>...</think>before the MMMU answer extractor) and--preserve_caller_max_tokens(keeps a larger--max_output_lengththan MMMU's defaultmax_gen_toks=512) for thinking models (e.g. Kimi K2.5, Step3p7) whose chain-of-thought breaks the default lm-eval regex.moe_op_backend.py): mapMiniMax2/SigmoidRenormrouting-method types, used by Step3p7's sigmoid router.Step3p7ForConditionalGenerationto the supported-models table and to both the model-feature and multimodal-feature support matrices, and update the speculative-decoding MTP note (MTP is no longer DeepSeek-only — it now covers any architecture shipping native MTP modules, including Step-3.x).Test Coverage
Unit tests
tests/unittest/_torch/modeling/test_modeling_step3p7.py— text decoder + MTP: stacked MoE weight splitting (per-expert key expansion, NVFP4 suffixes, no-op fallback), MTP weight rewriting to the MTP-block layout, multimodallanguage_model.*key flattening / prefix stripping, NVFP4 batched dequant round-trip, MTP-head pre-output normalization, AutoModel registration, plus real-checkpoint config/weight accounting, per-layer attention geometry helpers, and MTP spec-config defaulting to the checkpoint layer count. (Runs in L0 on A10.)tests/unittest/_torch/modeling/test_modeling_step3p7vl.py— vision tower + VLM wiring: RoPE-2D frequency cache / identity / dynamic-subgrid selection, LayerScale, fused-QKV vision attention layout (+ head-divisibility guard), encoder forward shape and positional-embedding interpolation, projector geometry and encode→text-hidden projection, weight-loading prefix splitting, VLM / vision-encoder / input-processor registration, vision-config geometry, and input-processor hooks (special and multimodal token ids, tokens-per-image, contiguous image span). (Runs in L0 on L40S.)tests/unittest/llmapi/test_llm_args.py::test_MTPDecodingConfig_default_draft_len_is_not_user_set— verifies an unsetmax_draft_lenstaysNone(and out ofmodel_fields_set), while an explicit value mirrors tomax_total_draft_tokens.Integration accuracy tests
tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestStep3_7— text-decoder GSM8K:test_auto_dtype(BF16, TP8/EP8),test_fp8_block_scales(TP4/EP4,mtp_nextn ∈ {0, 3}, KVCacheManagerV2, TRTLLM attention + MoE backends), andtest_nvfp4(Blackwell-only, TP4/EP4, FP8 KV cache). Reference scores intests/integration/defs/accuracy/references/gsm8k.yaml.tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py::TestStep3_7— multimodal MMMU:test_fp8_block_scalesandtest_nvfp4(Blackwell), using the strip-thinking post-processor +preserve_caller_max_tokens. Reference scores intests/integration/defs/accuracy/references/mmmu.yaml.CI registration
test_fp8_block_scales[...-mtp_nextn=3]) onl0_dgx_h100.yml; the two unit suites onl0_a10.yml(text) andl0_l40s.yml(vision).llm_function_core.txt): the remaining integration cases — text BF16 (TP8), FP8 without MTP, NVFP4, and both multimodal cases.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.