Skip to content

[TRTLLM-11950][perf] Audio feature extractor optimizations#14031

Open
2ez4bz wants to merge 1 commit into
NVIDIA:mainfrom
2ez4bz:dev-audio-preproc
Open

[TRTLLM-11950][perf] Audio feature extractor optimizations#14031
2ez4bz wants to merge 1 commit into
NVIDIA:mainfrom
2ez4bz:dev-audio-preproc

Conversation

@2ez4bz
Copy link
Copy Markdown
Collaborator

@2ez4bz 2ez4bz commented May 12, 2026

Summary by CodeRabbit

Release Notes

  • New Features

    • Added stereo audio support for audio processing pipelines.
    • Enhanced audio feature extraction with improved device compatibility and tensor handling.
    • Added sampling rate validation to catch audio format mismatches early.
  • Bug Fixes

    • Improved audio preprocessing robustness and consistency across model variants.

Review Change Stack

Description

This mirrors audio optimizations done in vLLM.

Perf impact

  • ISL=50, ISL=1k
  • Images = batch of 32 1536x1024 pixels
  • Audio = 30 min at 16KHz

H200 FP8

Scenario TTFT before / after (ms) Output TPS before / after
32 images, conc 1 2822.56 / 2779.03 119.07 / 119.28
32 images, conc 32 27,375.86 / 26,536.60 413.45 / 413.23
32 images, conc 64 29,974.22 / 29,795.18 435.40 / 436.35
     
30 min audio, conc 1 1,960.47 / 1,703.68 129.65 / 133.94
30 min audio, conc 32 19,320.61 / 16,672.18 838.14 / 866.00
30 min audio, conc 64 26,994.23 / 24,183.45 835.69 / 913.42

B200 NVFP4

Scenario TTFT before / after (ms) Output TPS before / after
32 images, conc 1 1,609.49 / 1,606.96 175.97 / 176.78
32 images, conc 32 14,658.01 / 14,233.86 752.37 / 762.44
32 images, conc 64 15,611.52 / 14,821.95 825.02 / 822.81
32 images, conc 128 26,989.92 / 34,641.86 855.21 / 825.25
     
30 min audio, conc 1 1,372.06 / 1,121.40 176.76 / 185.59
30 min audio, conc 32 16,074.67 / 16,213.79 1,115.21 / 1,217.24
30 min audio, conc 64 35,698.93 / 35,981.10 1,214.47 / 1,274.34

Test Coverage

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)

  • 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.

@2ez4bz 2ez4bz requested review from a team as code owners May 12, 2026 05:34
@2ez4bz 2ez4bz requested review from moraxu and yechank-nvidia May 12, 2026 05:34
@2ez4bz 2ez4bz changed the title [TRTLLM-11950] Audio feature extractor optimizations [TRTLLM-11950][perf] Audio feature extractor optimizations May 12, 2026
@2ez4bz
Copy link
Copy Markdown
Collaborator Author

2ez4bz commented May 12, 2026

/bot run

@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai Bot commented May 12, 2026

📝 Walkthrough

Walkthrough

ParakeetExtractor was refactored from a HuggingFace feature extractor wrapper to a Torch-native STFT-based log-mel feature extractor with cached signal processing components, device-aware tensor operations, and explicit argument validation. The extractor now returns a structured dict with audio features, attention masks, and clip counts. Nemotron Nano integration was simplified to use the new API directly, and comprehensive tests validate output structure, input flexibility, error handling, and configuration propagation.

Changes

ParakeetExtractor Torch Refactor and Nemotron Integration

Layer / File(s) Summary
ParakeetExtractor config and torch feature extraction
tensorrt_llm/_torch/models/modeling_parakeet.py
_ExtractorConfig extended with STFT parameters (hop_length, win_length, n_fft, preemphasis, padding_value) and from_hf_config constructor. ParakeetExtractor implementation refactored to torch-native STFT-to-log-mel feature computation with cached Hann windows, mel filter banks, optional preemphasis, and length-aware normalization producing attention masks.
ParakeetExtractor audio preprocessing and API
tensorrt_llm/_torch/models/modeling_parakeet.py
__call__ refactored to validate sampling_rate and return_tensors arguments, convert input audio to mono torch tensors, split into clips via split_audio_into_clips, apply optional preemphasis and normalization, and return dict with input_audio_features, feature_attention_mask, audio_num_clips on requested device. Token counting updated to use HF subsampling helper with frames tensor and minimum-1 guarantee.
Nemotron Nano integration with new extractor API
tensorrt_llm/_torch/models/modeling_nemotron_nano.py
NanoV2VLInputProcessor._prepare_audio_features simplified to call extractor(audios) directly and return expanded_text with extractor output dict, eliminating manual audio_data construction.
ParakeetExtractor behavior validation
tests/unittest/_torch/modeling/test_modeling_parakeet.py
Test suite updated to validate output keys (input_audio_features, feature_attention_mask, audio_num_clips), tensor shapes, stereo audio support, sampling rate validation raising ValueError, and config override propagation to both ext.config and extractor attributes.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~22 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 8.70% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check ❓ Inconclusive The PR description provides context (mirrors vLLM optimizations) and performance metrics showing measurable improvements, but the 'Test Coverage' section is empty, leaving what tests safeguard the changes unclear. Complete the 'Test Coverage' section to explicitly list the relevant tests (e.g., test_call_returns_expected_keys, test_call_accepts_stereo_audio, test_call_rejects_mismatched_sampling_rate, test_config_overrides) that validate the audio extractor changes.
✅ Passed checks (3 passed)
Check name Status Explanation
Title check ✅ Passed The title '[TRTLLM-11950] Audio feature extractor optimizations' directly and clearly summarizes the main change—refactoring audio preprocessing to optimize feature extraction performance.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
tensorrt_llm/_torch/models/modeling_parakeet.py (1)

44-49: 💤 Low value

Unnecessary int() wrapping on round() result.

In Python 3, round() already returns an int when called with a single argument (no decimal places specified). The int() call is redundant.

♻️ Suggested simplification
-        self._clip_target_samples = int(
-            round(self.config.clip_duration_s * self.config.sampling_rate)
-        )
-        self._tail_min_samples = int(
-            round(self.config.clip_min_duration_s * self.config.sampling_rate)
-        )
+        self._clip_target_samples = round(
+            self.config.clip_duration_s * self.config.sampling_rate
+        )
+        self._tail_min_samples = round(
+            self.config.clip_min_duration_s * self.config.sampling_rate
+        )
🤖 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 `@tensorrt_llm/_torch/models/modeling_parakeet.py` around lines 44 - 49, The
expressions assigning self._clip_target_samples and self._tail_min_samples wrap
round(...) with int(), which is redundant in Python 3; remove the outer int()
calls in the assignments to _clip_target_samples and _tail_min_samples (leave
round(self.config.clip_duration_s * self.config.sampling_rate) and
round(self.config.clip_min_duration_s * self.config.sampling_rate)) so the
result remains an int without the unnecessary cast.
🤖 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.

Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_parakeet.py`:
- Around line 44-49: The expressions assigning self._clip_target_samples and
self._tail_min_samples wrap round(...) with int(), which is redundant in Python
3; remove the outer int() calls in the assignments to _clip_target_samples and
_tail_min_samples (leave round(self.config.clip_duration_s *
self.config.sampling_rate) and round(self.config.clip_min_duration_s *
self.config.sampling_rate)) so the result remains an int without the unnecessary
cast.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 8f0c41e7-11ef-4ea8-8fde-19b204141865

📥 Commits

Reviewing files that changed from the base of the PR and between 7e07c8b and d5025b5.

📒 Files selected for processing (3)
  • tensorrt_llm/_torch/models/modeling_nemotron_nano.py
  • tensorrt_llm/_torch/models/modeling_parakeet.py
  • tests/unittest/_torch/modeling/test_modeling_parakeet.py

@tensorrt-cicd
Copy link
Copy Markdown
Collaborator

PR_Github #47885 [ run ] triggered by Bot. Commit: d5025b5 Link to invocation

Comment thread tensorrt_llm/_torch/models/modeling_parakeet.py
Comment thread tensorrt_llm/_torch/models/modeling_parakeet.py
Comment thread tensorrt_llm/_torch/models/modeling_parakeet.py Outdated
Comment thread tensorrt_llm/_torch/models/modeling_parakeet.py Outdated
Comment thread tensorrt_llm/_torch/models/modeling_parakeet.py Outdated
@2ez4bz 2ez4bz force-pushed the dev-audio-preproc branch from d5025b5 to dc137db Compare May 12, 2026 06:09
@2ez4bz 2ez4bz requested a review from a team as a code owner May 12, 2026 06:09
Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
@2ez4bz 2ez4bz force-pushed the dev-audio-preproc branch from dc137db to fffc64b Compare May 12, 2026 06:12
@2ez4bz
Copy link
Copy Markdown
Collaborator Author

2ez4bz commented May 12, 2026

/bot run

@tensorrt-cicd
Copy link
Copy Markdown
Collaborator

PR_Github #47900 [ run ] triggered by Bot. Commit: fffc64b Link to invocation

@tensorrt-cicd
Copy link
Copy Markdown
Collaborator

PR_Github #47900 [ run ] completed with state SUCCESS. Commit: fffc64b
/LLM/main/L0_MergeRequest_PR pipeline #37750 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants