fix: reduce memory occupation and raise performance for multi-nodes

[blueswhen]

No description provided.

[gemini-code-assist]

Code Review

This pull request introduces an optimized prefill MoE implementation using DeepEP's expanded dispatch and chunked reduction to dense rows, updating Dockerfiles, Triton kernels, and model layer inference files. Feedback focuses on improving robustness and compatibility: aligning intermediate byte sizes to 16-byte boundaries to prevent alignment crashes, relaxing shape checks in the workspace allocator, adding boundary masks in the accumulation kernels to support hidden sizes that are not multiples of 1024 (e.g., Qwen2-57B-A14B), and wrapping the private PyTorch allocator settings call in a hasattr check for backward compatibility.

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]

To prevent potential alignment crashes (e.g., RuntimeError when calling .view(torch.float32) on unaligned byte offsets) and to ensure optimal memory access performance on GPUs, all intermediate byte sizes and offsets should be aligned to at least 16-byte boundaries.

Currently, temp_row_bytes is calculated as the raw maximum of the three phases, which might not be a multiple of 16. Aligning temp_row_bytes to a 16-byte boundary guarantees that all subsequent workspace views (including qsilu_workspace and scale_workspace) are properly aligned.

Suggested change

	temp_row_bytes = max(
	silu_row_bytes + gemm_a_row_bytes,
	silu_row_bytes + quant_with_scale_row_bytes,
	gemm_b_row_bytes + quant_with_scale_row_bytes,
	)
	temp_row_bytes = max(
	silu_row_bytes + gemm_a_row_bytes,
	silu_row_bytes + quant_with_scale_row_bytes,
	gemm_b_row_bytes + quant_with_scale_row_bytes,
	)
	temp_row_bytes = (temp_row_bytes + 15) // 16 * 16

[gemini-code-assist]

The exact shape comparison tuple(shape) == scale_storage_shape is highly fragile. If per_token_group_quant_fp8 requests the unaligned shape (intermediate_size // block_size_k, chunk_rows) instead of the aligned shape (intermediate_size // block_size_k, aligned_chunk_rows), this check will fail and raise a RuntimeError, crashing the entire inference process.

To make this allocator robust, we should allow the requested column dimension to be less than or equal to the allocated aligned column dimension, and return a sliced view of scale_workspace matching the requested shape.

Suggested change

	def workspace_quant_alloc(shape, dtype, device):
	if tuple(shape) == tuple(qsilu_workspace.shape) and dtype == qsilu_workspace.dtype:
	return qsilu_workspace
	if tuple(shape) == scale_storage_shape and dtype == torch.float32:
	return scale_workspace
	raise RuntimeError(f"unexpected prefill quant allocation: shape={shape}, dtype={dtype}")
	def workspace_quant_alloc(shape, dtype, device):
	if tuple(shape) == tuple(qsilu_workspace.shape) and dtype == qsilu_workspace.dtype:
	return qsilu_workspace
	if dtype == torch.float32 and shape[0] == scale_storage_shape[0] and shape[1] <= scale_storage_shape[1]:
	return scale_workspace[:, :shape[1]]
	raise RuntimeError(f"unexpected prefill quant allocation: shape={shape}, dtype={dtype}")

[gemini-code-assist]

To support models with hidden sizes that are not multiples of 1024 (such as Qwen2-57B-A14B with hidden size 3584), we should accept hidden_size as a parameter and apply a boundary mask to all loads and stores.

Suggested change

	def _accumulate_expanded_chunk_kernel(
	total_recv_tokens,
	chunk,
	chunk_stride_m,
	chunk_stride_k,
	chunk_start,
	chunk_end,
	weights,
	recv_src_metadata,
	metadata_stride_m,
	metadata_stride_k,
	output,
	output_stride_m,
	output_stride_k,
	TOPK: tl.constexpr,
	BLOCK_D: tl.constexpr,
	):
	@triton.jit
	def _accumulate_expanded_chunk_kernel(
	total_recv_tokens,
	chunk,
	chunk_stride_m,
	chunk_stride_k,
	chunk_start,
	chunk_end,
	weights,
	recv_src_metadata,
	metadata_stride_m,
	metadata_stride_k,
	output,
	output_stride_m,
	output_stride_k,
	hidden_size,
	TOPK: tl.constexpr,
	BLOCK_D: tl.constexpr,
	):

[gemini-code-assist]

Apply a boundary mask mask = hidden_offsets < hidden_size to all loads and stores to safely handle hidden sizes that are not multiples of BLOCK_D (1024).

    hidden_block_id = tl.program_id(0)
    start_recv_token_id = tl.program_id(1)
    recv_token_grid_size = tl.num_programs(1)
    hidden_offsets = hidden_block_id * BLOCK_D + tl.arange(0, BLOCK_D)
    mask = hidden_offsets < hidden_size

    for recv_token_id in range(start_recv_token_id, total_recv_tokens, recv_token_grid_size):
        output_ptrs = output + recv_token_id * output_stride_m + hidden_offsets * output_stride_k
        accumulator = tl.load(output_ptrs, mask=mask, other=0.0).to(tl.float32)
        for topk_id in range(TOPK):
            slot = tl.load(
                recv_src_metadata
                + recv_token_id * metadata_stride_m
                + (topk_id + 2) * metadata_stride_k
            )
            if slot >= chunk_start and slot < chunk_end:
                local_row = (slot - chunk_start).to(tl.int64)
                value = tl.load(chunk + local_row * chunk_stride_m + hidden_offsets * chunk_stride_k, mask=mask, other=0.0)
                weight = tl.load(weights + slot)
                accumulator += value.to(tl.float32) * weight
        tl.store(output_ptrs, accumulator, mask=mask)

[gemini-code-assist]

The assertion assert chunk.shape[1] == output.shape[1] and output.shape[1] % block_d == 0 restricts the kernel to hidden sizes that are multiples of 1024. However, some popular MoE models (such as Qwen2-57B-A14B) have a hidden size of 3584, which is not a multiple of 1024, causing a hard crash on this assertion.

We should remove the divisibility assertion and pass output.shape[1] (the hidden size) to the Triton JIT kernel, using a boundary mask inside the kernel to safely handle any hidden size.

    """Accumulate one contiguous expanded W2 chunk into dense receive-token rows."""
    topk = recv_src_metadata.shape[1] - 2
    block_d = 1024
    assert chunk.shape[1] == output.shape[1]
    grid = (triton.cdiv(output.shape[1], block_d), min(output.shape[0], 1024))
    _accumulate_expanded_chunk_kernel[grid](
        output.shape[0],
        chunk,
        chunk.stride(0),
        chunk.stride(1),
        chunk_start,
        chunk_start + chunk.shape[0],
        weights,
        recv_src_metadata,
        recv_src_metadata.stride(0),
        recv_src_metadata.stride(1),
        output,
        output.stride(0),
        output.stride(1),
        output.shape[1],
        TOPK=topk,
        BLOCK_D=block_d,
        num_warps=2,
    )

[gemini-code-assist]

Using the private API torch._C._accelerator_setAllocatorSettings directly can cause an AttributeError and crash the application on startup if an older or different version of PyTorch (e.g., < 2.4) is used.

Wrapping this call in a hasattr check ensures backward compatibility and prevents startup crashes in environments with older PyTorch versions.

Suggested change

	torch._C._accelerator_setAllocatorSettings("expandable_segments:True")
	if hasattr(torch._C, "_accelerator_setAllocatorSettings"):
	torch._C._accelerator_setAllocatorSettings("expandable_segments:True")