Method test_chunked_prefill

tests/maxengine_test.py:165–298

Test identical result between chunked prefill with single and multiple chunked. The return value in kv_cache_prefill function is key and value itself. Although the value of key and value are the same as stored in the KVCache without quantization, the prefill still produce slightly diffe

(self)

  @pytest.mark.skip(reason="Can only pass on CPU.")
  def test_chunked_prefill(self):
    """Test identical result between chunked prefill with single and multiple chunked.

    The return value in kv_cache_prefill function is key and value itself.
    Although the value of key and value are the same as stored in the KVCache without quantization,
    the prefill still produce slightly different result while using multiple TPU devices or GPU due to unknown reasons.
    """

    prefill_length = 8
    tokens = jnp.array([1, 11, 22, 33, 444, 555, 666])
    padding_tokens = jnp.array([1, 11, 22, 33, 444, 555, 666, 0])
    true_length = tokens.shape[0]
    assert padding_tokens.shape[0] == prefill_length

    def array_equal_valid_tokens(x, y, *, compare_length):
      if len(x.shape) > 1:
        # containing sequence
        if x.shape[0] > 1:
          # Assume batch size is 1
          # sequence is the first axis for kv cache
          return jnp.array_equal(x[:compare_length], y[:compare_length])
        else:
          # sequence is the second axis for decoder segment id
          return jnp.array_equal(x[:, :compare_length], y[:, :compare_length])
      else:
        # single integer
        return jnp.array_equal(x, y)

    model_config_args = {
        "max_target_length": prefill_length * 4,
        "max_prefill_predict_length": prefill_length * 2,
        "model_call_mode": "inference",
        "capacity_factor": -1,
        "decoder_block": "mistral",
        "scan_layers": False,
        "per_device_batch_size": 1.0,
    }

    # Model without chunked prefill
    config = self.init_pyconfig(
        use_chunked_prefill=False,
        **model_config_args,
    )
    engine = MaxEngine(config)
    params = engine.load_params()
    expected_prefill_result, expected_first_token = engine.prefill(
        params=params,
        padded_tokens=padding_tokens,
        true_length=true_length,
    )

    # Model with chunked prefill
    config = self.init_pyconfig(
        use_chunked_prefill=True,
        **model_config_args,
    )
    engine = MaxEngine(config)
    params = engine.load_params()