deepseek v4 rebase track

[yueming-yuan]

Tracking draft PR for the DeepSeek V4 rebase onto Miles main.\n\nThis branch currently contains the full rebase track; follow-up work will split general Miles fixes/features into separate PRs so this PR can converge to DeepSeek V4 support only.

[gemini-code-assist]

Code Review

This pull request introduces full support for DeepSeek-V4 models, including Flash and Pro versions, by implementing specialized weight conversion, indexer replay, and a collection of TileLang kernels for optimized attention and hyper-connection operations. It also updates SGLang dependencies to v0.5.12 and adds a comprehensive training pipeline script. Technical feedback identifies several critical issues in the newly added kernels, such as missing boundary checks in the indexer and sparse MLA kernels that could cause out-of-bounds memory access. Furthermore, a bug in the FP8 simulation casting logic and a potential NameError in the weight update function were found, along with a recommendation to refactor model-specific padding detection for improved robustness.

[gemini-code-assist]

The variable num_new_engines is used here but it does not appear to be defined within the scope of the update_weights function in the current diff. This will result in a NameError at runtime. Please ensure it is correctly initialized, likely from the return value of a weight update or engine connection call.

[gemini-code-assist]

When inplace=True, out_dtype is reassigned to in_dtype (e.g., BF16). The inner T.Cast(out_dtype, ...) call then casts to BF16 instead of the intended FP8 quantization format. This prevents the kernel from correctly simulating FP8 quantization noise. You should use a separate variable to hold the quantization target type (e.g., the original FP8 constant).

                        y_local[i, j] = T.Cast(
                            out_dtype,
                            T.Cast(compute_dtype, T.Cast(FP8, T.clamp(
                                x_local[i, j] / s_local[i], fp8_min, fp8_max
                            ))) * s_local[i],
                        )

[gemini-code-assist]

The kernel writes to Logits without checking if the computed indices seq_len_i + bq_i and cu_k_s_min + nbn_i * block_N + bn_i are within the bounds of seq_len and seq_len_kv. This can lead to out-of-bounds memory access if the sequence lengths are not multiples of the block sizes.

Suggested change

	Logits[seq_len_i + bq_i, cu_k_s_min + nbn_i * block_N + bn_i] = logits[bn_i, bq_i]
	for bq_i, bn_i in T.Parallel(block_Q, block_N):
	if seq_len_i + bq_i < seq_len and cu_k_s_min + nbn_i * block_N + bn_i < seq_len_kv:
	Logits[seq_len_i + bq_i, cu_k_s_min + nbn_i * block_N + bn_i] = logits[bn_i, bq_i]

[gemini-code-assist]

The kernel performs atomic additions to dAttnSink using an index bz * block_H + h_i that can range up to padded_H - 1. If the number of heads H is not a power of 2 (e.g., H=43 or H=61), this will result in out-of-bounds memory access. A boundary check is required.

Suggested change

	T.atomic_add(
	dAttnSink[bz * block_H + h_i],
	-Delta[by, s_i, bz * block_H + h_i]
	* T.exp2(AttnSink[bz * block_H + h_i] * 1.44269504 - Lse[by, s_i, bz * block_H + h_i]),
	)
	for h_i in T.Parallel(block_H):
	if bz * block_H + h_i < H:
	T.atomic_add(
	dAttnSink[bz * block_H + h_i],
	-Delta[by, s_i, bz * block_H + h_i]
	* T.exp2(AttnSink[bz * block_H + h_i] * 1.44269504 - Lse[by, s_i, bz * block_H + h_i]),
	)

[gemini-code-assist]

The kernel accesses AttnSink using H0 + h_i, which can exceed the actual number of heads H because H_per_block is based on padded_H. This will cause out-of-bounds memory access for models where H is not a multiple of the block size. Please add a boundary check.

Suggested change

	for h_i in T.Parallel(H_per_block):
	sumexp[h_i] += T.exp2(AttnSink[H0 + h_i] * 1.44269504 - m_i[h_i] * sm_scale)
	for h_i in T.Parallel(H_per_block):
	if H0 + h_i < H:
	sumexp[h_i] += T.exp2(AttnSink[H0 + h_i] * 1.44269504 - m_i[h_i] * sm_scale)

[gemini-code-assist]

The logic for detecting when to adjust pad_size for DeepSeek-V4 models relies on checking dsv4_hc_mult. As noted in the TODO, this should be refactored to a more robust detection mechanism, such as checking the model type or a specific configuration flag, to avoid relying on a specific attribute that might be missing or zero in other valid configurations.

References

1. Model parameters should be retrieved from the model configuration rather than being hardcoded or inferred from potentially missing attributes.