CANN: fix RoPE cache issue on multi-device

[hipudding]

RoPE cache only needs to be computed once per token. However, in multi-device scenarios, not every device starts computation from layer 0, which may lead to unallocated memory issues and precision errors.

This commit records the first layer of each device to avoid the above issues.

Update
To avoid the RoPE cache being overly coupled to a specific model, we currently only cache those entries that can be determined, from the input parameters, not to undergo any transformation.

./bin/test-backend-ops test -b CANN0 -o ROPE
Testing 3 devices
Backend 1/3: CANN0
Device description: Ascend910B4
Device memory: 30196 MB (29802 MB free)
11837/11837 tests passed
Backend CANN0: OK
Backend 2/3: CANN1
Skipping
Backend 3/3: CPU
Skipping
3/3 backends passed
OK

Make sure to read the contributing guidelines before submitting a PR

[ggerganov]

This is really hacky. Can you improve without making assumptions about the tensor names? Maybe create the cache based on the input parameters?

[hipudding]

I’ve tried, but during the decode stage, it’s not possible to determine based on the shape and data of position, because all position lengths are the same, and position itself, as well as running position->data, are the same too. The only difference is the data inside position, but copying data from the device to the host is not a good approach. Do you have any good suggestions for this?

[hipudding]

I think I’ve come up with a method: during the forward computation in ggml_cgraph, add a marker when encountering the first RoPE operator and perform the cache calculation. The subsequent RoPE operators would then skip the computation. This way, we can avoid parsing the tensor’s name. I will try this way.

[noemotiovon]

The current scenario is that when computing the sine/cosine for the rope, it is calculated only for the first layer on each device, and other layers are reused. It is necessary to identify which layer is the first layer on the current device. However, currently there is no way to obtain the layer number within the device backend, so it can only be inferred from the name. Would it be possible to store the layer number information directly in the tensor?

[noemotiovon]

LGTM!

[noemotiovon]

This is an excellent refactor that removes the previous hacky implementation! We’ll do another refactor later for the case where src2 != nullptr.

[ggerganov]

AFAIU the cache will be computed for every graph_compute() call on the first rope operation and then will be reused for all remaining rope operations in the current graph. I think this assumes that all ropes are the same in all layers. In general this is not guaranteed and will likely cause problems in the future.

[hipudding]

AFAIU the cache will be computed for every graph_compute() call on the first rope operation and then will be reused for all remaining rope operations in the current graph. I think this assumes that all ropes are the same in all layers. In general this is not guaranteed and will likely cause problems in the future.

As far as I know, apart from freq_factors, the RoPE cache used for each token only depends on the position and some hyperparameters. So far, I haven’t observed any cases where the cache differs across layers. Could you clarify in what situations the RoPE cache would be layer-dependent? Thanks!

[ggerganov]

For example Gemma3n uses different freq_base for the different layers:

llama.cpp/src/llama-model.cpp

Lines 10635 to 10639 in f15d515

	for (int il = 0; il < n_layer; ++il) {
	// this block is made to be closely resemble Gemma3p5DecoderLayer on python code
	const float freq_base_l = model.get_rope_freq_base (cparams, il);
	const float freq_scale_l = model.get_rope_freq_scale(cparams, il);

Also, this simple test program would likely not run correctly if it was offloaded to the CANN backend (currently it always runs on the CPU):

llama.cpp/tests/test-rope.cpp

Lines 156 to 177 in f15d515

	if (m < 3) {
	struct ggml_tensor * p0 = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ne[2]);
	struct ggml_tensor * p1 = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ne[2]);
	struct ggml_tensor * p2 = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ne[2]);
	for (int i = 0; i < ne[2]; ++i) {
	((int32_t *) p0->data)[i] = n_past_0 + i;
	((int32_t *) p1->data)[i] = n_past_2 - n_past_0;
	((int32_t *) p2->data)[i] = n_past_2 + i;
	}
	// test mode 0, 2, 4 (standard, GPT-NeoX, GLM)
	mode = m == 0 ? 0 : m == 1 ? 2 : 4;
	// 100, 101, 102, ..., 172
	r0 = ggml_rope(ctx0, x, p0, n_rot, mode);
	// -67, -67, -67, ..., -67
	r1 = ggml_rope(ctx0, r0, p1, n_rot, mode); // "context swap", i.e. forget n_past_0 - n_past_2 tokens
	// 33, 34, 35, ..., 105
	r2 = ggml_rope(ctx0, x, p2, n_rot, mode);
	} else {
	// testing multi-dimension rope position embedding mode

Generally, we want to keep the code generic and not make assumptions about the application. llama.cpp is just one of the applications that use ggml and other applications are possible to call ggml_rope with different parameters within the same graph.

[hipudding]

@ggerganov Thank you for the reminder. All unsafe caches have been removed, and only the parts that can be determined through parameters to remain unchanged are cached.

[noemotiovon]

LGTM. However, for transformer models, removing the sin/cos cache in ROPE leads to a performance drop compared to before. We’ll need to explore a more elegant way to determine positional information in the future to ensure the cache check remains precise.