[float8] read dim from kwargs in Float8Tensor aten.split handler

[GotFusion]

Fixes #4430.

Summary

The aten.split.Tensor handler on Float8Tensor (added in #3334) unpacks three positionals from args:

tensor, split_size_or_sections, dim = args

That works for Tensor.chunk(n, dim=d) and other call sites where the dispatcher hands dim over in args. It does not work for the free-function form torch.chunk(t, n, dim=d), which reaches the handler with args = (tensor, split_size_or_sections) and dim in kwargs:

ValueError: not enough values to unpack (expected 3, got 2)

This breaks every fully_shard(module) call where the module's weights have been quantized via quantize_(model, Float8DynamicActivationFloat8WeightConfig(...)), because FSDP2's _chunk_with_empty is implemented as:

chunks = list(torch.chunk(tensor, num_chunks, dim=dim))

Minimal reproducer

import torch
import torchao  # noqa: F401
from torch.distributed._composable.fsdp import fully_shard
from torchao.quantization import (
    Float8DynamicActivationFloat8WeightConfig,
    quantize_,
)
from torchao.quantization.granularity import PerRow
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-0.5B-Instruct", torch_dtype=torch.bfloat16
).to("cuda")
quantize_(
    model,
    Float8DynamicActivationFloat8WeightConfig(granularity=PerRow()),
    device="cuda",
)
for block in model.model.layers:
    fully_shard(block)
fully_shard(model)
# -> ValueError: not enough values to unpack (expected 3, got 2)

Originally hit under FSDP2 + 4-card NPU (torchao_npu); same code path triggers on CUDA + FSDP2 too.

Fix

Read dim from either positional or keyword form, matching the aten.split.Tensor schema (Tensor self, SymInt split_size, int dim=0):

tensor, split_size_or_sections = args[0], args[1]
dim = args[2] if len(args) > 2 else kwargs.get("dim", 0)

Test Plan

Adds test_chunk_via_torch_chunk_with_dim_kwarg next to the existing test_chunk. The existing test only exercises Tensor.chunk(...) (the vLLM Llama 4 path from #3334), which is why this regression was not caught. The new test calls torch.chunk(t, n, dim=dim) explicitly:

pytest test/quantization/quantize_/workflows/float8/test_float8_tensor.py -s -x -k 'chunk'

The new test fails before the handler fix with ValueError: not enough values to unpack (expected 3, got 2), and passes after.

Related

• Issue: [float8] Float8Tensor aten.split.Tensor handler crashes on torch.chunk(t, n, dim=<kwarg>) (breaks FSDP2 fully_shard) #4430
• Handler was introduced in add support for torch.chunk to float8tensor #3334 (LLaMA 4 MoE weight loading via vLLM).
• This PR only changes the argument extraction; the body of the handler is untouched.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4429

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Run pull request jobs on OSDC runners in shadow mode

⚠️ 1 Awaiting Approval

As of commit 9a21bc8 with merge base b506da5 ():

AWAITING APPROVAL - The following workflow needs approval before CI can run:

• PR Label Check (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vkuzo]

@claude review this PR

[claude]

Claude finished @vkuzo's task in 52s —— View job

---

PR Review: [float8] read dim from kwargs in Float8Tensor aten.split handler

• Read changed files and understand the fix
• Review the implementation for correctness
• Review the test coverage
• Post review feedback

---

Overall: Looks good. This is a clean, minimal fix for a real bug that affects FSDP2 + Float8 quantization.

Implementation (torchao/quantization/quantize_/workflows/float8/float8_tensor.py:1029-1030)

The fix correctly handles both calling conventions:

tensor, split_size_or_sections = args[0], args[1]
dim = args[2] if len(args) > 2 else kwargs.get("dim", 0)

• The default dim=0 matches the ATen schema: aten.split.Tensor(Tensor self, SymInt split_size, int dim=0) — correct.
• The rest of the handler body is untouched — good, keeps the diff minimal.

Test (test/quantization/quantize_/workflows/float8/test_float8_tensor.py:1440-1452)

The new test_chunk_via_torch_chunk_with_dim_kwarg test directly exercises the failing path (torch.chunk(t, n, dim=dim) which passes dim as a kwarg). The split-then-cat-then-compare approach is a clean round-trip correctness check.

Minor notes (non-blocking)

1. The test tensor shape (16, 5120, 16384) is quite large (~2.5 GB in bf16). This matches the existing test_chunk test so presumably CI can handle it, but worth keeping in mind if memory-constrained runners are added in the future.

2. The comment on lines 1027-1028 is fine — it explains the why (different callers pass dim differently), which is exactly the kind of non-obvious behavior worth documenting.

No blocking issues found. LGTM.

---