[Model Runner V2] Introduce num_tokens_for_attn

[WoosukKwon]

No description provided.

[mergify]

Hi @WoosukKwon, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[gemini-code-assist]

Code Review

This pull request introduces a new field, num_tokens_for_attn, to BatchExecutionDescriptor and InputBatch to allow for a more precise specification of the number of tokens used in attention mechanisms, particularly for CUDA graphs. The changes are consistently applied across cudagraph_utils, dp_utils, model_runner, and model_states to propagate and utilize this new field. The implementation appears correct and well-integrated.

[gemini-code-assist]

Code Review

This pull request introduces a new field, num_tokens_for_attn, to BatchExecutionDescriptor and InputBatch. This field is used to specify the exact number of tokens that should be processed by the attention mechanism, which can differ from the total number of tokens in a batch, particularly when using CUDA graphs for decode operations. The changes are consistently propagated through cudagraph_utils.py, dp_utils.py, input_batch.py, model_runner.py, and model_states/default.py. This refactoring centralizes the logic for determining the attention token count and simplifies the prepare_attn function. The implementation appears correct and robust, with no high or critical issues found.

[njhill]

we probably wanna rename these: num_actual_tokens <= num_attn_tokens <= num_input_tokens

agree and let's document the above relationship explicitly in comment too :)

[njhill]

save a line?

Suggested change

	# i.e. no request padding is needed
	# so we leave it as None
	# i.e. no request padding is needed, so we leave it as None

[njhill]

I think we should add more doc to these fields

including meaning of None for the other ones too

[njhill]

could simplify

Suggested change

	if batch_desc.num_tokens_for_attn is not None:
	num_tokens_for_attn = batch_desc.num_tokens_for_attn
	else:
	num_tokens_for_attn = num_tokens
	num_tokens_for_attn = batch_desc.num_tokens_for_attn or num_tokens

[njhill]

Also would be good to make sure the CI tests cover this