amd/deepseek_v4 integration 1/N - 0426

[HaiShaw]

Motivation

Update amd/deepseek_v4 integration branch

Following PRs have large set of conflict, we use this PR and upstream amd/deepseek_v4 branch to integrate in parallel.
#23600
#23608

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

[gemini-code-assist]

Code Review

This pull request introduces support for DeepSeek-V4, featuring compressed attention backends, a specialized indexer, and speculative decoding (NextN) with optimized JIT kernels. Key updates include a redesigned memory pool and fused metadata initialization. Feedback identifies critical logic and performance issues in the DeepseekV4BackendRadix and DeepseekV4MultiStepBackend implementations, specifically addressing off-by-one errors in speculative decoding loops, performance bottlenecks from host-device synchronization, and platform compatibility concerns on AMD/ROCm systems.

[gemini-code-assist]

On AMD/ROCm platforms (is_hip()), the flash_mla kernel is typically not available. Defaulting to "kernel" will lead to an import error or crash. It is better to default to "torch" when running on HIP to ensure the fallback path is used unless explicitly overridden.

Suggested change

	import os
	backend = os.environ.get("SGLANG_HACK_FLASHMLA_BACKEND", "torch")

[gemini-code-assist]

The loop range range(self.speculative_num_steps - 1) will skip the last backend in self.attn_backends. It should be range(self.speculative_num_steps) to ensure all backends are initialized.

Suggested change

	for i in range(self.speculative_num_steps - 1):
	for i in range(self.speculative_num_steps):

[gemini-code-assist]

The loop range range(1, self.speculative_num_steps - 1) skips the last backend in self.attn_backends. For example, if speculative_num_steps is 2, the range is empty and the second backend is never updated. It should be range(1, self.speculative_num_steps). Additionally, self.forward_metadata for the MultiStepBackend instance itself should be set to temp_metadata to avoid using stale or None metadata during the forward pass.

        temp_metadata = self.attn_backends[0].forward_metadata
        self.forward_metadata = temp_metadata

        # Copy to other backends without recomputing
        for i in range(1, self.speculative_num_steps):

[gemini-code-assist]

Calling .item() on a GPU tensor causes a host-device synchronization, which can significantly degrade performance in the hot path. Additionally, if num_tokens is 0, req_pool_indices_repeated[-1] will raise an IndexError. Consider using torch.cat with expand to perform the padding entirely on the device without synchronization.

            if num_tokens > 0:
                padding = req_pool_indices_repeated[-1:].expand(pad_size)
            else:
                padding = req_pool_indices_repeated.new_zeros(pad_size)
            req_pool_indices_repeated = torch.cat([req_pool_indices_repeated, padding])

[gemini-code-assist]

This function performs a torch.cat even when no padding is required (i.e., when size == tensor.shape[0]), which is inefficient. Also, it will fail if size < tensor.shape[0]. Adding a check to return the original tensor when padding is not needed would be better.

    if size <= tensor.shape[0]:
        return tensor
    if value == 0:
        return torch.cat(
            [tensor, tensor.new_zeros(size - tensor.shape[0], *tensor.shape[1:])],
            dim=0,
        )