Dynamic MTP

[shihaobai]

No description provided.

[gemini-code-assist]

Code Review

This pull request introduces support for the Gemma-4 model, including multimodal vision processing, custom prefill attention with image bidirectional masking, and logit softcapping. It also refactors memory management by introducing KvCacheAllocator and decoupled operators, adds SymmMem and FlashInfer all-reduce fast paths, supports SM100 GPUs with FP4/FP8 Mega MoE, and implements an experimental Anthropic Messages API compatibility layer. Feedback on the changes highlights several critical issues: a runtime TypeError in prefill_cuda_graph.py from using set[int] as a callable, potential AttributeErrors in basemodel.py and grouped_fused_moe_ep.py due to missing attributes or uninitialized buffers, and an ignored stride parameter in the gather_token_id.py Triton kernel.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Using set[int] as a callable at runtime will raise a TypeError: 'GenericAlias' object is not callable in Python 3.9+ (or a TypeError: 'type' object is not subscriptable in older versions). It should be changed to the standard set constructor.

Suggested change

	graph_handle_token_nums = list(set[int](graph_handle_token_nums))
	graph_handle_token_nums = list(set(graph_handle_token_nums))

[gemini-code-assist]

The ReqManager class does not define a req_sampling_params_manager attribute. Accessing self.req_manager.req_sampling_params_manager.req_to_next_token_ids will raise an AttributeError. You should access req_to_next_token_ids directly from self.req_manager.

Suggested change

	req_to_next_token_ids=self.req_manager.req_sampling_params_manager.req_to_next_token_ids,
	req_to_next_token_ids=self.req_manager.req_to_next_token_ids,

[gemini-code-assist]

If is_sm100_gpu() is True but use_sm100_mega_moe(quant_method) is False (for example, when running FP8 MoE on an SM100 GPU), dist_group_manager.ep_low_latency_buffer remains None because it is only initialized when not is_sm100_gpu(). This will cause an AttributeError during decode. Consider ensuring that ep_low_latency_buffer is initialized or falling back to ep_buffer when ep_low_latency_buffer is None.

[gemini-code-assist]

The stride parameter req_to_next_token_ids_stride_1 is passed to the kernel but ignored inside the load instruction, effectively hardcoding the stride of the second dimension to 1. If req_to_next_token_ids is non-contiguous, this will result in incorrect memory access.

Suggested change

	cur_next_token_id = tl.load(
	req_to_next_token_ids + cur_req_idx * req_to_next_token_ids_stride + cur_mtp_index, mask=block_mask
	)
	cur_next_token_id = tl.load(
	req_to_next_token_ids + cur_req_idx * req_to_next_token_ids_stride + cur_mtp_index * req_to_next_token_ids_stride_1, mask=block_mask
	)