JK just uses AJP to communicate with Tomcat - and is not aware of HTTP session clustering or JBoss partitions. Define your JK groups/channels to match your Tomcat partitions/instances.
Partition A has nodes A1 and A2
Partition B has nodes B1 and B2
In JK configure a group each for partitions A and B. Then define channels for A1, A2 and make them a member of group A. Likewise for B.
Map application for partition A to group A, by its URI. Likewise for B.
Note: If you have multiple Tomcats on a single box, you will need to define distinct AJP ports for each Tomcat. And make sure your JK config (the channels portion) is pointing at the right AJP ports.