[Question] IndexShare: why is the retained indexer supervised only by its own layer, not the multi-layer objective from the IndexCache paper?

### Question

In the DSA IndexShare support added by #5099, only the **full** (source) layer computes an
indexer loss; **shared** layers contribute no loss term at all. The IndexCache paper that
IndexShare comes from specifies a multi-layer distillation objective instead, where the
retained indexer is distilled against *every* layer it serves. I would like to understand
whether the single-layer objective is a deliberate simplification, and if so, what motivated
it.

### What the implementation does

`experimental/lite/megatron/lite/primitive/modules/attention/dsa.py`, both the native and the
cuDNN-fused path, gate the indexer loss on `not self.skip_topk`:

```python
# L1034 (native) and L1268 (cuDNN fused), identical condition
if self.training and torch.is_grad_enabled() and not self.skip_topk:
indexer_loss = _cp_indexer_loss(
q_indexer, k_indexer, weights_indexer, topk_indices,
query.detach(), kv.detach(), ...)
out = DSAIndexerLossAutoScaler.apply(out, indexer_loss)
```

with `self.skip_topk = indexer_type == "shared"` (L769). A shared layer also builds no indexer
module at all (`self.indexer: DSAIndexer | None = None`, L799), so it has no projections of its
own to feed a loss with.

So the retained indexer is supervised only by the distribution of the layer it lives in, and
the layers that reuse its top-k never enter the objective.

### What the paper specifies

*IndexCache: Accelerating Sparse Attention via Cross-Layer Index Reuse*
([arXiv:2603.12201](https://arxiv.org/abs/2603.12201)) defines the multi-layer distillation
loss as

$$\mathcal{L}_{\text{multi}}^{I} = \sum_{j=0}^{m}\frac{1}{m+1}\sum_{t} D_{KL}\big(p_t^{(\ell+j)}\,\|\,q_t^{(\ell)}\big)$$

where $\ell$ is the retained full layer and $\ell+1 \dots \ell+m$ are the shared layers it
serves. The KL targets $p_t^{(\ell+j)}$ are **each served layer's own attention distribution**,
averaged over the group; the student $q_t^{(\ell)}$ is the single retained indexer. The paper's
Proposition 1 states $\nabla_\theta \mathcal{L}_{\text{multi}}^{I} = \nabla_\theta
\mathcal{L}_{\text{avg}}^{I}$, i.e. it is gradient-equivalent to distilling against the averaged
target distribution, so that the indexer learns a top-k selection covering the tokens that
matter across all the layers it serves.

The paper applies this objective in both training phases:

> "In the warm-up phase, we train the indexer in the F layer using $\mathcal{L}_{\text{multi}}^{I}$,
> while keeping all other parameters fixed. In the sparse training phase, we continue to train
> the indexer using $\mathcal{L}_{\text{multi}}^{I}$ ... and additionally include the LM loss"

### Questions

1. Is supervising the retained indexer on its own layer only a deliberate choice, or does it
simply fall out of this implementation path (shared layers hold no indexer, so there are no
projections to build the extra terms from)?
2. If deliberate — was the accuracy impact of dropping the served layers' targets measured?
With larger groups the retained indexer is applied to more layers whose distributions it
never saw.
3. Is a multi-layer objective planned? Implementing it here seems to need either shared layers
calling the source layer's indexer module with their own q/k/softmax statistics, or a fused
kernel that accepts a pre-averaged teacher — the current
`SparseLightningIndexerGradKLLoss`-style kernels rebuild the teacher from a single layer's
statistics and have no slot for a pre-averaged one, which makes the multi-layer objective
cost one kernel call per served layer rather than one per group.

Context: we are implementing the same feature in MindSpore MindFormers and currently follow the
paper's multi-layer objective. Aligning with the reference implementation matters to us, so
understanding the reasoning behind this difference would be very helpful. Thanks!

Start in experimental/lite/megatron/lite/primitive/modules/attention/dsa.py, especially the indexer setup and the native and cuDNN-fused loss gates, then compare them with the IndexCache paper's multi-layer objective. Done means the intended behavior, accuracy evidence, and scope for any multi-layer objective are clarified or an implementation plan is agreed.

Written by the indexing model from the issue text.