fix: fix guided decoding state corruption in turbomind when tp>1

[windreamer]

Motivation

When serving models with tensor parallelism (--tp >= 2), enabling guided decoding (e.g., response_format: {"type": "json_schema"}) causes segmentation faults。

Root Cause: The guided decoding state management layers (GuidedDecodeMaskLayer and GuidedDecodeUpdateLayer) were being instantiated on all tensor parallelism (TP) ranks. Each rank independently modified shared decoding state, causing race conditions and memory corruption. In distributed inference, only rank 0 should orchestrate guided decoding logic while other ranks perform parallel computation.

Modification

Eliminate shared mutable state by allocating an independent GrammarMatcher instance per TP rank. Enhance the decoding pipeline with rank awareness so each thread operates on its own isolated matcher, ensuring thread safety while maintaining deterministic behavior across ranks.

Fixes #4152

[irexyc]

The sampled token are not broadcasted from tp_rank0, so all ranks should do the same sampling process to make sure the next token is same on all ranks.

I think the problem may be the state of GrammarMatcher can not be shared by all ranks. Currently, all ranks share the same std::shared_ptr<xgrammar::GrammarMatcher> in the request.

[windreamer]

GrammarMatcher

OK so we need to copy the GrammarMatcher instead ?

[windreamer]

The sampled token are not broadcasted from tp_rank0, so all ranks should do the same sampling process to make sure the next token is same on all ranks.

I think the problem may be the state of GrammarMatcher can not be shared by all ranks. Currently, all ranks share the same std::shared_ptr<xgrammar::GrammarMatcher> in the request.

I believe that a quick fix should be making GuidedDecodeUpdateLayer only executed in rank 0. As we only modify GrammarMatcher here. But I have no idea how to ensure we call GuidedDecodeUpdateLayer::Forward only when all ranks finish there sampling.

[irexyc]

I believe that a quick fix should be making GuidedDecodeUpdateLayer only executed in rank 0.

All ranks should have the same next token and the next token are computed by dynamic decoding, so we should make sure the dynamic decoding process are same on all ranks.

I think we can construct n_ranks of matchers here like r->matchers = ... https://github.com/InternLM/lmdeploy/blob/main/src/turbomind/engine/model_request.cc#L131C9-L131C19

and chose the correspond matcher here like matchers_.push_back(r->matchers[rank_]);
https://github.com/InternLM/lmdeploy/blob/main/src/turbomind/layers/sampling_layers/GuidedDecodeMaskLayer.cc#L36
https://github.com/InternLM/lmdeploy/blob/main/src/turbomind/layers/sampling_layers/GuidedDecodeUpdateLayer.cc#L32

[lzhangzz]

1. Pass h_tp_group to GuidedDecodeUpdateLayer
2. h_tp_group->Sync() here so that all ranks completed filling their host mask buffer before any rank tries to update matcher state.
3. AcceptToken when h_tp_group->rank() == 0

In addition, a stream sync is required after the copy. need_apply as in GuidedDecodeMaskLayer is neede here to avoid the copy / sync cost when guided decoding is not needed.

[windreamer]

Sadly, not only AcceptToken but also FillNextTokenBitmask modify the shared matcher state. So if we need multiple times of sync and shared bit mask, that will kill the performance I believe.

So I take @irexyc 's advice to just dup the state for each thread.

[lvhan028]

Vote for making GuidedDecodeUpdateLayer only executed in rank 0 and calling GuidedDecodeUpdateLayer::Forward only when all ranks finish the sampling.
It doesn't make sense to me that ModelRequest holds a field tp_size_.