[#15182][fix] Fix embedding vocab mask for handling rejection sampling in Kimi-K2.5

[chungen04]

Summary by CodeRabbit

• Bug Fixes
  • Improved embedding operation handling to correctly apply masking and remapping logic across different tensor parallel configurations.

Description

Fix #15182

Quick recap of the issue One-model Eagle3 speculative decoding with use_rejection_sampling: true crashes during the first generation step (warmup during engine launch) with a device-side assert, where the issue is specific to Kimi (not to the models supported as indicated in #12588):

  /opt/pytorch/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1515:
  indexSelectSmallIndex: ... Assertion `srcIndex < srcSelectDimSize` failed.

Summary of cause The root cause (examined by logging) is that the draft model's token embedding receives flashinfer's -1 "rejected token" vocab id, and on models such as Kimi-K2.6, the embedding does not vocab-mask its input, so F.embedding(-1) faults. Models such as Llama-4-Maverick, running correctly with use_rejection_sampling: true absorb the same -1 via the existing vocab-mask and never crash.

Code tracing of the root cause

1. Where the -1 comes from: this is why the problem is specific to use_rejection_sampling: true.

Rejection acceptance calls flashinfer's chain_speculative_sampling via rejection_sampling_one_model (tensorrt_llm/_torch/speculative/one_model_sampler.py). Its returned output_token_ids pads rejected positions with -1 (flashinfer sampling.py, chain_speculative_sampling docstring: "rejected samples are padded with -1").

This path is taken in _sample_and_accept_draft_tokens_rejection (tensorrt_llm/_torch/speculative/interface.py), and the raw accepted_tokens (still containing -1) is returned unchanged.

2. How the -1 reaches the draft embedding

prepare_1st_drafter_inputs builds the draft model's input_ids by flattening all accepted-token slots, including the rejected (-1) ones src:

input_ids_gen = accepted_tokens[num_contexts:,:spec_metadata.runtime_draft_len + 1].flatten()

That tensor flows into the draft loop (eagle3.py _run_draft_forward), which calls the draft token embedding (modeling_speculative.py:414,Eagle3DraftModel.forward):

inputs_embeds = self.embed_tokens(input_ids).to(self.dtype)

The assert listed in the original crashing log is in tensorrt_llm/_torch/modules/embedding.py, pre_comm_embedding_ops:

output = F.embedding(input_, weight)

input_ got -1, which cause the kernel to got caught in assertion.

3. Why this is specific to Kimi / DeepSeek-V3, and why Llama-4-Maverick does not break

tl;dr: Llama guarded the -1 vocab ids via the vocab id mapping to TP shards, which is enabled by default, but not for other models such as DeepSeek-V3.

The case for Llama-4-Maverick: the tracing stack is as follows.

• With use_rejection_sampling, the -1 vocab ids were emitted as usual.
• Looking around the embedding call embedding.py#L172, embedding.py#L164 triggers for llama-4, as llama-4 embedding layer is initialized with tp_mode src.
• The get_masked_input_and_mask method indicates that the -1 vocab id is not mapped to any shard in TP ranks. output is eventually mask filled with 0 src. This path naturally handles out of range IDs.

For Kimi / DeepSeek-V3:

• Same, with use_rejection_sampling, the -1 vocab ids were emitted upon rejection.
• The embed_tokens in DeepSeek-V3 is not initialized with TP modeling_deepseekv3.py#L1759. The -1 vocab id was directly used to index the embedding table, causing the crash.

Proposed fix

Derived from the trace, there are a few possible fixes:

• Enable TP in embed_tokens in class DeepseekV3Model.
• In tensorrt_llm/_torch/modules/embedding.py, pre_comm_embedding_ops: apply the existing get_masked_input_and_mask in the replicated (tp_mode is None) case as well, using the local vocab range [0, weight.shape[0]).

The second was what current PR adopts. This is a minimal change that reuses the same masking primitive Llama-4 already relies on, and it does not touch the rejection sampler, the draft loop, or accepted_tokens (where -1 remains a meaningful vocab id that downstream code reads via num_accepted_tokens, which is not altered here).

Test Coverage

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR extends the input token masking logic in the embedding operation to support non-tensor-parallel mode. The pre_comm_embedding_ops function now conditionally masks and remaps token IDs before embedding lookup, covering both tensor-parallel and non-parallel configurations to prevent invalid indices from reaching the embedding table.

Changes

Embedding Masking for Non-Tensor-Parallel Mode

Layer / File(s)	Summary
Embedding masking extension for non-tensor-parallel mode tensorrt_llm/_torch/modules/embedding.py	The pre_comm_embedding_ops function now initializes input_mask = None and extends the masking/remapping logic beyond TensorParallelMode.COLUMN to also handle tp_mode is None using the full embedding weight vocabulary range, ensuring subsequent output masking/padding correctly depends on whether an input_mask was produced.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The code changes directly address the bug reported in #15182 by fixing embedding vocab masking logic that was causing out-of-bounds indexing during rejection sampling in the Kimi path.
Out of Scope Changes check	✅ Passed	All changes in tensorrt_llm/_torch/modules/embedding.py are narrowly scoped to fixing the embedding vocab mask issue, directly addressing #15182 requirements without extraneous modifications.
Description check	✅ Passed	The PR description is comprehensive, detailing the root cause of the issue with clear code tracing, explaining why it affects Kimi but not Llama-4, and justifying the proposed fix.
Title check	✅ Passed	The title clearly and specifically identifies the main change: fixing the embedding vocab mask for handling rejection sampling in Kimi-K2.5, directly aligned with the PR objectives and code changes.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[zhaoyangwang-nvidia]

Thanks for the fix — this correctly closes the crash on the non-TP (attention-DP) path. Two small follow-ups:

1. Add a comment on where the -1 comes from. It's not obvious why masking is needed here. flashinfer's chain_speculative_sampling pads non-accepted tokens with -1, and these invalid ids reach the embedding lookup — that's what this guards against. Suggested:

elif tp_mode is None:
    # flashinfer's rejection kernel (chain_speculative_sampling) pads non-accepted
    # tokens with -1. In non-TP mode the full vocab is local, so mask out-of-range
    # ids (e.g. -1) over [0, weight.shape[0]) to avoid an OOB embedding lookup.
    input_, input_mask = get_masked_input_and_mask(input_, 0, weight.shape[0])

2. Could you add the same guard for ROW mode? ROW also holds the full vocab on each rank, so it has the same -1 exposure but isn't masked. Note the ROW all-gather padding is in an elif that's mutually exclusive with the mask-cleanup, so it needs to be decoupled:

input_mask = None
if tp_mode == TensorParallelMode.COLUMN:
    input_, input_mask = get_masked_input_and_mask(input_, vocab_start_index, vocab_end_index)
elif tp_mode is None or tp_mode == TensorParallelMode.ROW:
    # Full vocab present locally -> mask out-of-range ids (e.g. -1).
    input_, input_mask = get_masked_input_and_mask(input_, 0, weight.shape[0])

output = F.embedding(input_, weight)

if input_mask is not None:
    output.masked_fill_(input_mask, 0)

# ROW: pad hidden dim for all-gather (independent of masking).
if tp_mode == TensorParallelMode.ROW and gather_output:
    if tp_rank == tp_size - 1 and padding_size > 0:
        output = F.pad(output, (0, padding_size))

return output

[chungen04]

Thank you for the feedback. I am taking @zhaoyangwang-nvidia 's recommendation here as it also guards >= vocab_size vocab id (weight.shape[0]) and reuses the helper function (get_masked_input_and_mask). Let me know if there are more comments.

[yuxianq]

Thank you for the feedback. I am taking @zhaoyangwang-nvidia 's recommendation here as it also guards >= vocab_size vocab id (weight.shape[0]) and reuses the helper function (get_masked_input_and_mask). Let me know if there are more comments.

It is Ok to reuse get_masked_input_and_mask, but we still don't handle the row TP case for rejection sampling, right? To fix it, we need to replace elif tp_mode is None: with else to cover the row TP case and call output.masked_fill_(input_mask, 0) in all cases.

[chungen04]

@yuxianq resolved. I was thinking about making the condition of else explicit, but it's fine since the comment also explains the situation.

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #53817 [ run ] triggered by Bot. Commit: e3bf30f Link to invocation

[tensorrt-cicd]

PR_Github #53817 [ run ] completed with state SUCCESS. Commit: e3bf30f
/LLM/main/L0_MergeRequest_PR pipeline #42932 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #53867 [ run ] triggered by Bot. Commit: e3bf30f Link to invocation

[tensorrt-cicd]

PR_Github #53867 [ run ] completed with state SUCCESS. Commit: e3bf30f
/LLM/main/L0_MergeRequest_PR pipeline #42974 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #53872 [ run ] triggered by Bot. Commit: e3bf30f Link to invocation

[tensorrt-cicd]

PR_Github #53872 [ run ] completed with state SUCCESS. Commit: e3bf30f
/LLM/main/L0_MergeRequest_PR pipeline #42977 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54026 [ run ] triggered by Bot. Commit: 5771281 Link to invocation

[tensorrt-cicd]

PR_Github #54026 [ run ] completed with state SUCCESS. Commit: 5771281
/LLM/main/L0_MergeRequest_PR pipeline #43110 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[zhaoyangwang-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #54295 [ run ] triggered by Bot. Commit: 5771281 Link to invocation

[tensorrt-cicd]

PR_Github #54295 [ run ] completed with state SUCCESS. Commit: 5771281
/LLM/main/L0_MergeRequest_PR pipeline #43366 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[chungen04]

Should we rebase the branch?

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54425 [ run ] triggered by Bot. Commit: 660b505 Link to invocation

[tensorrt-cicd]

PR_Github #54425 [ run ] completed with state FAILURE. Commit: 660b505
/LLM/main/L0_MergeRequest_PR pipeline #43494 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[xxi-nv]

/bot run

[tensorrt-cicd]

PR_Github #54528 [ run ] triggered by Bot. Commit: 660b505 Link to invocation

[tensorrt-cicd]

PR_Github #54528 [ run ] completed with state SUCCESS. Commit: 660b505
/LLM/main/L0_MergeRequest_PR pipeline #43585 completed with status: 'SUCCESS'

CI Report

Link to invocation

[zhaoyangwang-nvidia]

Hi @yuxianq this PR is ready to merge, please help to merge it.