feat: add json object constrained decoding. (#2110)

Co-authored-by: shifengmin <shifengmin@qq.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: 3 people

docs/src/content/docs/en/features/mtp.md

@@ -88,6 +88,40 @@ Input model references:
 
 When using MTP for inference, you need to specify both the main model and the draft model (MTP model).
 
+## JSON object output
+
+The OpenAI-compatible Chat Completions API supports token-level JSON-object
+constrained decoding for models whose tokenizer can provide stable token
+pieces. Use:
+
+```json
+{
+  "response_format": {"type": "json_object"}
+}
+```
+
+This mode supports ordinary generation and MTP, including streaming and
+non-streaming responses. With PD separation enabled, the prefill instance
+forwards the format flag and reasoning mode to the decode instance. The
+decode sequence commits and validates the first token produced by prefill
+before building its next-token mask.
+
+The feature is model-independent at the runtime level. A request is rejected
+if the model tokenizer cannot provide the token pieces required to build the
+grammar. For models with reasoning enabled, the tokenizer must also encode the
+existing `</think>` boundary so the runtime can begin JSON enforcement after
+reasoning completes.
+
+When reasoning is enabled, reasoning remains unconstrained until the model's
+`</think>` boundary; JSON enforcement starts afterward. When thinking is
+disabled through the existing chat-template kwargs, enforcement starts at the
+first generated token.
+
+The current MVP supports only `json_object`. `json_schema`, regex/custom
+grammars, legacy Completion requests, C API request structs, and tool-call
+structural tags are not supported. Unsupported `response_format.type` values
+are rejected as invalid requests.
+
 #### DeepSeek-V3/V3.2/R1 Launch Example
 ```bash
 MODEL_PATH="/models/DeepSeek-V3"

docs/src/content/docs/zh/features/mtp.md

@@ -86,6 +86,34 @@ python3 tools/export_mtp.py \
 
 使用MTP进行推理时，需要同时指定主模型和草稿模型（MTP模型）。
 
+## JSON 对象输出
+
+OpenAI 兼容的 Chat Completions 接口支持基于 tokenizer token piece 的通用
+token 级 JSON 对象约束。请求中设置：
+
+```json
+{
+  "response_format": {"type": "json_object"}
+}
+```
+
+该模式支持普通生成和 MTP，也支持流式与非流式返回。开启 PD 分离时，prefill
+实例会把格式标志和 reasoning 状态传递给 decode 实例；decode 侧会先提交并校验
+prefill 产出的第一个 token，再生成后续 token 的约束 mask。
+
+该能力在 runtime 层不依赖具体模型类型。如果模型 tokenizer 无法提供构建
+grammar 所需的稳定 token piece，请求会被拒绝。对于启用 reasoning 的模型，
+tokenizer 还必须能够编码已有的 `</think>` 边界，runtime 才能在 reasoning
+结束后开始 JSON 约束。
+
+开启 reasoning 时，`</think>` 标记之前的 reasoning token 不受 JSON 约束，标记
+之后才开始约束；通过现有 chat-template kwargs 关闭 thinking 时，从第一个生成
+token 开始约束。
+
+当前 MVP 只支持 `json_object`。暂不支持 `json_schema`、正则/自定义 grammar、
+legacy Completion、C API 请求结构体以及 tool-call 结构化标签。其他
+`response_format.type` 会作为非法请求拒绝。
+
 #### DeepSeek-V3/V3.2/R1 启动示例
 ```bash
 MODEL_PATH="/models/DeepSeek-V3"

setup.py

@@ -691,6 +691,8 @@ def finalize_options(self) -> None:
         super().finalize_options()
 
     def run(self) -> None:
+        global BUILD_TEST_FILE
+
         build_ext_cmd = self.get_finalized_command("build_ext")
         build_ext_cmd.device = self.device
         build_ext_cmd.arch = self.arch
@@ -719,7 +721,6 @@ def run(self) -> None:
                 path = os.path.join(root, item)
                 if "_test" in item and os.path.isfile(path):
                     os.remove(path)
-        global BUILD_TEST_FILE
         BUILD_TEST_FILE = False
 
         self.skip_build = True

tests/core/framework/batch/batch_packed_input_test.cpp

@@ -26,6 +26,7 @@ limitations under the License.
 #include "core/framework/block/block_manager_impl.h"
 #include "core/framework/model/model_input_params.h"
 #include "core/framework/request/stopping_checker.h"
+#include "core/framework/sampling/json_object_grammar.h"
 #include "core/runtime/forward_params.h"
 #include "core/runtime/params_utils.h"
 
@@ -184,6 +185,9 @@ TEST(BatchPackedInputTest, PackedProtoLazyUnpackRestoresSampleIdxes) {
       builder.build_forward_input(/*num_decoding_tokens=*/1,
                                   /*min_decoding_batch_size=*/0);
   ASSERT_TRUE(input.sampling_params.sample_idxes.defined());
+  input.sampling_params.filter_bitmask =
+      torch::tensor({{static_cast<int32_t>(0x5)}},
+                    torch::TensorOptions().dtype(torch::kInt32));
 
   proto::PackedForwardInput packed_input;
   ASSERT_TRUE(forward_input_to_packed_proto(input, &packed_input));
@@ -203,6 +207,76 @@ TEST(BatchPackedInputTest, PackedProtoLazyUnpackRestoresSampleIdxes) {
   ASSERT_TRUE(unpacked_input.sampling_params.sample_idxes.defined());
   EXPECT_TRUE(tensor_equals_vector<int32_t>(
       unpacked_input.sampling_params.sample_idxes, {0}));
+  ASSERT_TRUE(unpacked_input.sampling_params.filter_bitmask.defined());
+  EXPECT_TRUE(torch::equal(unpacked_input.sampling_params.filter_bitmask,
+                           input.sampling_params.filter_bitmask));
+}
+
+TEST(BatchPackedInputTest, PackedProtoLazyToPreservesJsonMetadata) {
+  RequestSamplingParam sampling_param;
+  StoppingChecker stopping_checker;
+  stopping_checker.set_max_generated_tokens(4);
+
+  SequenceParams seq_params;
+  seq_params.seq_capacity = 32;
+  seq_params.stopping_checker = &stopping_checker;
+  seq_params.sampling_param = &sampling_param;
+  seq_params.enable_schedule_overlap = true;
+
+  MMData mm_data;
+  BlockManager::Options options;
+  options.num_blocks(2).block_size(4);
+  BlockManagerImpl manager(options);
+
+  IncrementalDecoder decoder("", 1, false, false);
+  Sequence sequence(/*index=*/0,
+                    /*token_ids=*/{1, 2, 3, 4},
+                    torch::Tensor(),
+                    mm_data,
+                    std::move(decoder),
+                    seq_params);
+  sequence.add_blocks(BlockType::KV, manager.allocate(1));
+
+  std::vector<Sequence*> sequences = {&sequence};
+  std::vector<uint32_t> allowed_max_tokens = {4};
+  BatchInputBuilder builder(sequences,
+                            allowed_max_tokens,
+                            {},
+                            {},
+                            nullptr,
+                            /*batch_id=*/2,
+                            nullptr,
+                            BatchForwardType::DECODE);
+  ForwardInput input = builder.build_forward_input(
+      /*num_decoding_tokens=*/1, /*min_decoding_batch_size=*/0);
+
+  JsonObjectGrammar grammar({"{", "}", "stop"}, /*stop_token_ids=*/{2});
+  JsonObjectGrammarState state = grammar.initial_state();
+  ASSERT_TRUE(state.accept_token(/*open_object=*/0));
+  input.json_object_state_snapshots = {state.snapshot()};
+  input.sample_sequence_ids = {"req-json#0"};
+  input.sample_prior_output_rows = {-1};
+
+  proto::PackedForwardInput packed_input;
+  ASSERT_TRUE(forward_input_to_packed_proto(input, &packed_input));
+
+  ForwardInput lazy_input;
+  packed_proto_to_forward_input(
+      packed_input, lazy_input, torch::Device(torch::kCPU), nullptr);
+  EXPECT_TRUE(lazy_input.input_host_buffer_has_layout);
+  EXPECT_TRUE(lazy_input.json_object_state_snapshots.empty());
+
+  const ForwardInput materialized_input =
+      lazy_input.to(torch::Device(torch::kCPU), torch::kFloat32);
+  ASSERT_EQ(materialized_input.json_object_state_snapshots.size(), 1u);
+  EXPECT_EQ(materialized_input.json_object_state_snapshots[0].token_ids,
+            std::vector<int32_t>({0}));
+  EXPECT_FALSE(
+      materialized_input.json_object_state_snapshots[0].reasoning_enabled);
+  EXPECT_EQ(materialized_input.sample_sequence_ids,
+            std::vector<std::string>({"req-json#0"}));
+  EXPECT_EQ(materialized_input.sample_prior_output_rows,
+            std::vector<int32_t>({-1}));
 }
 
 }  // namespace xllm