Skip to content

DsaTopkKernel sizes chosen[512]/picked[64] by literal while index_topk is 512 (Flash) / 1024 (Pro) #505

Description

@localai-bot

DsaTopkKernel (src/vt/cuda/cuda_deepseek_v4.cu:624-665) sizes two thread-local arrays by literal:

bool chosen[512];   // :639  indexed [0, n) where n = candidate-window length
int64_t picked[64]; // :641  written [0, topk)

Neither bound is asserted, and neither is derived from the config. topk is the caller's index_topk, which is 512 for V4-Flash and 1024 for V4-Pro — both far past picked[64]. The overflow branch is the n > topk path, so it is entered whenever the candidate window is larger than index_topk.

DispTopk (src/vllm/model_executor/models/deepseek_v4.cpp:336-340) routes to this kernel whenever be.device is set, with no size guard.

Reachability — narrower than it first looks

This is not currently reachable on the shipped real-model path. dsa_dense = (be.gguf != nullptr) (deepseek_v4.cpp:668) forces is_indexer false on the keep-quant GGUF path, so the real Flash run never calls the indexer at all; the kernel is exercised only at the collapsed synthetic geometry (be.gguf == nullptr), where topk is small by construction and the comment at :639 ("nk small in the structural gate") holds.

So this is a latent trap in a test-only path, not a shipped defect — filed because the real-geometry DSA sparse path is a named residual, and the moment that residual is built out these two literals become a silent thread-stack overflow at the real index_topk rather than a loud failure. The host reference DsaTopkSelect (deepseek_v4_dsa.cpp:72) has no such bound; it allocates from n.

Upstream has no equivalent constraint: selection goes through ops.top_k_per_row_prefill (vllm/model_executor/layers/sparse_attn_indexer.py:488-497), and the candidate window is built as ks = row_start, ke = row_start + (pos + 1) // COMPRESS_RATIO (vllm/v1/attention/backends/mla/indexer.py:270-290) — i.e. the full causal prefix in compressed-key space, with no fixed cap.

Fix

Assert both bounds against the launch parameters so the kernel refuses rather than corrupts, and size the selection from the config instead of a literal when the real-geometry path lands. The minimal change is the guard plus a test that drives the device topk with topk > 64 and confirms it is rejected rather than silently wrong.

Found while assessing #504.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions