[Perf] Add VLLM_TRITON_FORCE_FIRST_CONFIG to skip Triton autotuning

[fuscof-ibm]

Purpose

Triton's @triton.autotune benchmarks every candidate config on the first launch per key, then caches the winner. Three things make this painful when debugging or measuring:

1. Autotune is timing-driven, so its winner is non-deterministic. Run-to-run jitter can promote a different (BLOCK_M, num_warps, ...) tuple each time, and a different tuple reduces partial sums in a different order — so identical inputs produce numerically different outputs across runs.

2. Cached results bleed across investigations. Triton's autotune cache is persisted on disk, so "just rerun it" doesn't reset state. Editing every kernel's configs=[...] list to pin a value, or wiping caches between runs, is fragile and error-prone (plus third-party libs ship their own autotuned kernels we don't own).

3. No visibility into what was picked. If autotuning chose a different config on a particular GPU/run, you only find out by digging.

This PR adds a debug-only knob, VLLM_TRITON_FORCE_FIRST_CONFIG, that short-circuits autotune at runtime by walking the candidate configs in declaration order and using the first one that compiles+launches without OutOfResources / CompileTimeAssertionFailure / PTXASError. The picked index is cached per (autotuner, key) so subsequent calls stay deterministic and cheap. One INFO line per unique kernel is logged so you can see exactly which config ran.

Default off — no behavior change unless explicitly enabled.

Sample output

With VLLM_TRITON_FORCE_FIRST_CONFIG=1:

[force_first_config.py:91] [triton-autotune-disabled] kernel=chunk_scaled_dot_kkt_fwd_kernel configs=27 picked_index=0 picked=BK: 32, num_warps: 2, num_ctas: 1, num_stages: 2, maxnreg: None
[force_first_config.py:91] [triton-autotune-disabled] kernel=chunk_gated_delta_rule_fwd_kernel_h_blockdim64 configs=12 picked_index=0 picked=BV: 32, num_warps: 2, num_ctas: 1, num_stages: 2, maxnreg: None
[force_first_config.py:91] [triton-autotune-disabled] kernel=chunk_fwd_kernel_o configs=24 picked_index=0 picked=BK: 64, BV: 64, num_warps: 2, num_ctas: 1, num_stages: 2, maxnreg: None

Relationship to #34648

#34648 (VLLM_TRITON_AUTOTUNE) targets the same root cause but with a different mechanism and a different default behavior. Summary of differences:

	#34648	This PR
Default	Disables autotune by default (changes vLLM behavior)	Off — pure opt-in debug knob
Patch site	triton.autotune (the decorator) — passes configs[:1]	Autotuner.run (the runtime method)
Reaches autotuners already constructed at import time (e.g. in fla, flashinfer)	Only if disable_autotune_globally() runs before those modules are imported (import-order dependent)	Yes — patches the class method, takes effect on next launch
Bad-config handling	Crashes if configs[0] OOMs or fails to compile on the target GPU	Walks past invalid configs, picks first valid one
Per-kernel visibility	None	One INFO line per unique kernel (count + picked config)
Batch-invariance integration	Auto-enabled under VLLM_BATCH_INVARIANT=1	None (independent debug knob)

The two have different purposes: #34648 is a user-facing toggle intended to make determinism the default; this PR is a runtime-level patch intended for debugging and perf measurement, with no default behavior change. The runtime-method approach is what allowed it to intercept third-party autotuners we don't decorate ourselves.

How this was used to debug #40172

While preparing #40172 (fused Triton kernel for GPU-side Mamba state postprocessing in hybrid models, now merged), we observed (under MTP) a different number of accepted tokens across runs that didn't correlate with code changes (see #40172 (comment)) . Re-running the same configuration produced different number of accepted tokens, particularly in the GDN prefill path.

Setting VLLM_TRITON_FORCE_FIRST_CONFIG=1 for both baseline and patched runs eliminated that noise.

This is the workflow that motivated the patch: a one-line env-var flip that makes Triton autotune behave deterministically for everything in the process, without needing to know in advance which kernels are involved or owning the code that decorates them.

Test Plan

Unit tests covering the patched Autotuner.run logic with a fake autotuner instance (CPU-only, no GPU required as long as triton is importable):

.venv/bin/python -m pytest tests/test_force_first_config.py -vvv

The tests cover:

• skip-invalid-configs[0] and pick the next valid config
• per-(autotuner, key) caching: invalid config is not retried on
  subsequent calls
• empty configs falls back to direct call
• all-invalid raises RuntimeError with the original OutOfResources
  chained as __cause__

End-to-end smoke test: vllm serve on a GDN-using model with the env var set, confirming the per-kernel INFO lines appear and the server serves requests:

VLLM_TRITON_FORCE_FIRST_CONFIG=1 vllm serve "Qwen/Qwen3.5-9B" \
    --port 8007 -tp 1 -pp 1 -dp 1 \
    --language-model-only \
    --gdn-prefill-backend=triton \
    --speculative-config '{"method":"qwen3_next_mtp","num_speculative_tokens":2}'

Test Result

tests/test_force_first_config.py::test_skips_invalid_first_config_and_caches_second PASSED
tests/test_force_first_config.py::test_empty_configs_falls_back_to_direct_call PASSED
tests/test_force_first_config.py::test_all_configs_invalid_raises_runtime_error PASSED

3 passed

Pre-commit hooks (ruff-check, ruff-format, mypy-3.12, SPDX, typos) pass on the changed files.

Duplicate-work check (per AGENTS.md)

• gh pr list --search "triton autotune in:title,body" — only related open
  PR is [Feature] Add VLLM_TRITON_AUTOTUNE with functional autotune control #34648, addressed above; [Core] Add shmem-aware autotune pruner for non-H100 Triton kernels #43047 (shmem-aware autotune pruner) is
  orthogonal (config pruning for low-shmem GPUs, not determinism).
• gh pr list --search "VLLM_TRITON in:title,body" — no other matches.
  .

[gemini-code-assist]

Code Review

This pull request introduces a mechanism to disable Triton's autotuning by monkey-patching the Autotuner.run method to always select the first configuration when the VLLM_TRITON_FORCE_FIRST_CONFIG environment variable is enabled. Feedback indicates that the implementation should handle cases where self.configs is empty to avoid an IndexError and should call self.fn directly instead of self.fn.run to ensure compatibility with Triton 3.x.

[gemini-code-assist]

The code assumes self.configs is non-empty. If a kernel is defined with an empty list of configurations (which is allowed in Triton), this will raise an IndexError. The original Triton Autotuner.run implementation handles this by checking if self.configs is empty and falling back to a direct call.

Suggested change

	config = self.configs[0]
	if not self.configs:
	return self.fn(*args, **kwargs)
	config = self.configs[0]

[fuscof-ibm]

The current code already has the if statement https://github.com/fuscof-ibm/vllm/blob/1e443852b201ab596ebdcd77eb1ad3d772339908/vllm/triton_utils/force_first_config.py#L38

[gemini-code-assist]

Using self.fn.run(...) will cause an AttributeError on Triton 3.x when the autotuner wraps a Heuristics object, as Heuristics does not have a run method in newer Triton versions. It is safer and more compatible to call self.fn(...) directly, which is what the upstream Triton Autotuner.run does.

Suggested change

	return self.fn.run(*args, **kwargs, **config.all_kwargs())
	return self.fn(*args, **kwargs, **config.all_kwargs())

[fuscof-ibm]

Even newer Triton versions have the run method (see https://github.com/triton-lang/triton/blob/7eef28877ea964619aeb0c52d1e521d087819b61/python/triton/runtime/autotuner.py#L474).

We have anyway we fall back to self.fn in case the run is not available (see https://github.com/fuscof-ibm/vllm/blob/1e443852b201ab596ebdcd77eb1ad3d772339908/vllm/triton_utils/force_first_config.py#L64)

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @fuscof-ibm.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[tdoublep]

LGTM

This is very useful for debugging performance-related issue when using Triton kernels.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @fuscof-ibm.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork